[mypyc] Use table-driven helper for from imports

Add CPyImport_ImportFromMany() which imports a module and a tuple of
names, placing them in the globals dict in one go.

Previously, each name would imported and placed in globals manually
in IR, leading to some pretty verbose code.

The other option to collect all from imports and perform them all at
once in the helper would remove even more ops, however, it has some
major downsides:

- It wouldn't be able to be done in IRBuild directly, instead being
  handled in the prebuild visitor and codegen... which all sounds
  really involved.

- It would cause from imports to be performed eagerly, potentially
  causing circular imports (especially in functions whose imports are
  probably there to avoid a circular import!).

The latter is the nail in the coffin for this idea.

mypyc/irbuild/builder.py

@@ -89,7 +89,6 @@
     RType,
     RUnion,
     bitmap_rprimitive,
-    c_int_rprimitive,
     c_pyssize_t_rprimitive,
     dict_rprimitive,
     int_rprimitive,
@@ -126,12 +125,7 @@
 from mypyc.primitives.dict_ops import dict_get_item_op, dict_set_item_op
 from mypyc.primitives.generic_ops import iter_op, next_op, py_setattr_op
 from mypyc.primitives.list_ops import list_get_item_unsafe_op, list_pop_last, to_list
-from mypyc.primitives.misc_ops import (
-    check_unpack_count_op,
-    get_module_dict_op,
-    import_extra_args_op,
-    import_op,
-)
+from mypyc.primitives.misc_ops import check_unpack_count_op, get_module_dict_op, import_op
 from mypyc.primitives.registry import CFunctionDescription, function_ops
 
 # These int binary operations can borrow their operands safely, since the
@@ -394,22 +388,6 @@ def add_to_non_ext_dict(
         key_unicode = self.load_str(key)
         self.call_c(dict_set_item_op, [non_ext.dict, key_unicode, val], line)
 
-    def gen_import_from(
-        self, id: str, globals_dict: Value, imported: list[str], line: int
-    ) -> Value:
-        self.imports[id] = None
-
-        null_dict = Integer(0, dict_rprimitive, line)
-        names_to_import = self.new_list_op([self.load_str(name) for name in imported], line)
-        zero_int = Integer(0, c_int_rprimitive, line)
-        value = self.call_c(
-            import_extra_args_op,
-            [self.load_str(id), globals_dict, null_dict, names_to_import, zero_int],
-            line,
-        )
-        self.add(InitStatic(value, id, namespace=NAMESPACE_MODULE))
-        return value
-
     def gen_import(self, id: str, line: int) -> None:
         self.imports[id] = None
 

mypyc/irbuild/statement.py

@@ -43,13 +43,16 @@
     YieldFromExpr,
 )
 from mypyc.ir.ops import (
+    NAMESPACE_MODULE,
     NO_TRACEBACK_LINE_NO,
     Assign,
     BasicBlock,
     Branch,
+    InitStatic,
     Integer,
     LoadAddress,
     LoadErrorValue,
+    LoadLiteral,
     MethodCall,
     RaiseStandardError,
     Register,
@@ -96,7 +99,7 @@
 from mypyc.primitives.misc_ops import (
     check_stop_op,
     coro_op,
-    import_from_op,
+    import_from_many_op,
     send_op,
     type_op,
     yield_from_except_op,
@@ -258,29 +261,25 @@ def transform_import_from(builder: IRBuilder, node: ImportFrom) -> None:
         module_package = ""
 
     id = importlib.util.resolve_name("." * node.relative + node.id, module_package)
-
-    globals = builder.load_globals_dict()
-    imported_names = [name for name, _ in node.names]
-    module = builder.gen_import_from(id, globals, imported_names, node.line)
-
-    # Copy everything into our module's dict.
+    builder.imports[id] = None
+
+    names = [name for name, _ in node.names]
+    as_names = [as_name or name for name, as_name in node.names]
+    names_literal = builder.add(LoadLiteral(tuple(names), object_rprimitive))
+    if names == as_names:
+        # Reuse names tuple to reduce verbosity.
+        as_names_literal = names_literal
+    else:
+        as_names_literal = builder.add(LoadLiteral(tuple(as_names), object_rprimitive))
     # Note that we miscompile import from inside of functions here,
-    # since that case *shouldn't* load it into the globals dict.
+    # since that case *shouldn't* load everything into the globals dict.
     # This probably doesn't matter much and the code runs basically right.
-    for name, maybe_as_name in node.names:
-        as_name = maybe_as_name or name
-        obj = builder.call_c(
-            import_from_op,
-            [module, builder.load_str(id), builder.load_str(name), builder.load_str(as_name)],
-            node.line,
-        )
-        builder.gen_method_call(
-            globals,
-            "__setitem__",
-            [builder.load_str(as_name), obj],
-            result_type=None,
-            line=node.line,
-        )
+    module = builder.call_c(
+        import_from_many_op,
+        [builder.load_str(id), names_literal, as_names_literal, builder.load_globals_dict()],
+        node.line,
+    )
+    builder.add(InitStatic(module, id, namespace=NAMESPACE_MODULE))
 
 
 def transform_import_all(builder: IRBuilder, node: ImportAll) -> None:

mypyc/lib-rt/CPy.h

@@ -622,8 +622,8 @@ PyObject *CPy_Super(PyObject *builtins, PyObject *self);
 PyObject *CPy_CallReverseOpMethod(PyObject *left, PyObject *right, const char *op,
                                   _Py_Identifier *method);
 
-PyObject *CPyImport_ImportFrom(PyObject *module, PyObject *package_name,
-                               PyObject *import_name, PyObject *as_name);
+PyObject *CPyImport_ImportFromMany(PyObject *mod_id, PyObject *names, PyObject *as_names,
+                                   PyObject *globals);
 
 PyObject *CPySingledispatch_RegisterFunction(PyObject *singledispatch_func, PyObject *cls,
                                              PyObject *func);

mypyc/lib-rt/misc_ops.c

@@ -670,8 +670,8 @@ CPy_Super(PyObject *builtins, PyObject *self) {
 }
 
 // This helper function is a simplification of cpython/ceval.c/import_from()
-PyObject *CPyImport_ImportFrom(PyObject *module, PyObject *package_name,
-                               PyObject *import_name, PyObject *as_name) {
+static PyObject *CPyImport_ImportFrom(PyObject *module, PyObject *package_name,
+                                      PyObject *import_name, PyObject *as_name) {
     // check if the imported module has an attribute by that name
     PyObject *x = PyObject_GetAttr(module, import_name);
     if (x == NULL) {
@@ -702,6 +702,31 @@ PyObject *CPyImport_ImportFrom(PyObject *module, PyObject *package_name,
     return NULL;
 }
 
+PyObject *CPyImport_ImportFromMany(PyObject *mod_id, PyObject *names, PyObject *as_names,
+                                   PyObject *globals) {
+    PyObject *mod = PyImport_ImportModuleLevelObject(mod_id, globals, 0, names, 0);
+    if (mod == NULL) {
+        return NULL;
+    }
+
+    for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(names); i++) {
+        PyObject *name = PyTuple_GET_ITEM(names, i);
+        PyObject *as_name = PyTuple_GET_ITEM(as_names, i);
+        PyObject *obj = CPyImport_ImportFrom(mod, mod_id, name, as_name);
+        if (obj == NULL) {
+            Py_DECREF(mod);
+            return NULL;
+        }
+        int ret = CPyDict_SetItem(globals, as_name, obj);
+        Py_DECREF(obj);
+        if (ret < 0) {
+            Py_DECREF(mod);
+            return NULL;
+        }
+    }
+    return mod;
+}
+
 // From CPython
 static PyObject *
 CPy_BinopTypeError(PyObject *left, PyObject *right, const char *op) {

mypyc/primitives/misc_ops.py

@@ -10,7 +10,6 @@
     c_pyssize_t_rprimitive,
     dict_rprimitive,
     int_rprimitive,
-    list_rprimitive,
     object_pointer_rprimitive,
     object_rprimitive,
     str_rprimitive,
@@ -120,25 +119,11 @@
     error_kind=ERR_MAGIC,
 )
 
-# Import with extra arguments (used in from import handling)
-import_extra_args_op = custom_op(
-    arg_types=[
-        str_rprimitive,
-        dict_rprimitive,
-        dict_rprimitive,
-        list_rprimitive,
-        c_int_rprimitive,
-    ],
+# From import helper op
+import_from_many_op = custom_op(
+    arg_types=[object_rprimitive, object_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,
-    c_function_name="PyImport_ImportModuleLevelObject",
-    error_kind=ERR_MAGIC,
-)
-
-# Import-from helper op
-import_from_op = custom_op(
-    arg_types=[object_rprimitive, str_rprimitive, str_rprimitive, str_rprimitive],
-    return_type=object_rprimitive,
-    c_function_name="CPyImport_ImportFrom",
+    c_function_name="CPyImport_ImportFromMany",
     error_kind=ERR_MAGIC,
 )