[mypyc] Getting capsule pointer from module instead of PyCapsule_Import (#18286)

advait-dixit · web-flow · commit f445369482e0 · 2024-12-19T18:48:03.000Z
Fixes mypyc/mypyc#999. `PyCapsule_Import` was failing in sub-packages. Since the capsule is an attribute of the module, we can access the capsule from the module instead of importing it.
diff --git a/mypyc/lib-rt/module_shim.tmpl b/mypyc/lib-rt/module_shim.tmpl
@@ -5,8 +5,11 @@ PyInit_{modname}(void)
 {{
     PyObject *tmp;
     if (!(tmp = PyImport_ImportModule("{libname}"))) return NULL;
+    PyObject *capsule = PyObject_GetAttrString(tmp, "init_{full_modname}");
     Py_DECREF(tmp);
-    void *init_func = PyCapsule_Import("{libname}.init_{full_modname}", 0);
+    if (capsule == NULL) return NULL;
+    void *init_func = PyCapsule_GetPointer(capsule, "{libname}.init_{full_modname}");
+    Py_DECREF(capsule);
     if (!init_func) {{
         return NULL;
     }}
diff --git a/mypyc/test-data/commandline.test b/mypyc/test-data/commandline.test
@@ -243,3 +243,22 @@ def i(arg: Foo) -> None:
 
 [file test.py]
 names = (str(v) for v in [1, 2, 3])  # W: Treating generator comprehension as list
+
+[case testSubPackage]
+# cmd: pkg/sub/foo.py
+from pkg.sub import foo
+
+[file pkg/__init__.py]
+
+[file pkg/sub/__init__.py]
+print("importing...")
+from . import foo
+print("done")
+
+[file pkg/sub/foo.py]
+print("imported foo")
+
+[out]
+importing...
+imported foo
+done
