Add bucketing test

Author: JukkaL

mypyc/test-data/run-multimodule.test

@@ -113,6 +113,28 @@ root_dir = os.path.dirname(native.__file__)
 assert os.path.dirname(native.other_sub.__file__) == os.path.join(root_dir, "other_pkg"), \
     native.other_sub.__file__
 
+[case testImportFromBucketingEndToEnd]
+# separate: [(["native.py", "other_pkg/__init__.py", "other_pkg/other_sub.py"], "testgroup")]
+from other_pkg import other_sub as sub_alias, other_attr as attr_alias, py_sub as py_alias
+[file other_pkg/__init__.py]
+other_attr = 41
+[file other_pkg/other_sub.py]
+def func() -> int:
+    return 10
+[file other_pkg/py_sub.py]
+def func() -> int:
+    return 20
+[file driver.py]
+import importlib
+import native
+other_sub = importlib.import_module("other_pkg.other_sub")
+py_sub = importlib.import_module("other_pkg.py_sub")
+assert native.sub_alias is other_sub
+assert native.attr_alias == 41
+assert native.py_alias is py_sub
+assert native.sub_alias.func() == 10
+assert native.py_alias.func() == 20
+
 [case testNativeTopLevelModuleFileDir]
 # separate: [(["native.py", "other_top.py"], "testgroup")]
 import other_top

mypyc/test/test_statement.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import unittest
+
+from mypy.nodes import Import, MypyFile
+from mypy.options import Options
+from mypyc.errors import Errors
+from mypyc.irbuild.builder import IRBuilder, IRVisitor
+from mypyc.irbuild.mapper import Mapper
+from mypyc.irbuild.prebuildvisitor import PreBuildVisitor
+from mypyc.irbuild.statement import (
+    IMPORT_NATIVE_ATTR,
+    IMPORT_NATIVE_SUBMODULE,
+    IMPORT_NON_NATIVE,
+    classify_import_from,
+    group_consecutive,
+    import_globals_id_and_name,
+    split_import_group_to_python_and_native,
+)
+from mypyc.options import CompilerOptions
+
+
+class DummyVisitor(IRVisitor):
+    pass
+
+
+def make_builder(
+    *,
+    module_name: str = "pkg.current",
+    native_modules: set[str] | None = None,
+    same_group_modules: set[str] | None = None,
+    graph: set[str] | None = None,
+) -> IRBuilder:
+    native_modules = native_modules or set()
+    same_group_modules = same_group_modules or set()
+    group_map: dict[str, str | None] = {module_name: "current-group"}
+    for module in native_modules:
+        group_map[module] = "current-group" if module in same_group_modules else "other-group"
+
+    errors = Errors(Options())
+    current_file = MypyFile([], [])
+    current_file._fullname = module_name
+    pbv = PreBuildVisitor(errors, current_file, {}, {})
+    builder = IRBuilder(
+        module_name,
+        {},
+        {name: object() for name in (graph or set())},
+        errors,
+        Mapper(group_map),
+        pbv,
+        DummyVisitor(),
+        CompilerOptions(),
+        {},
+    )
+    builder.set_module(module_name, module_name.replace(".", "/") + ".py")
+    return builder
+
+
+class TestStatementHelpers(unittest.TestCase):
+    def test_import_globals_id_and_name_for_plain_import(self) -> None:
+        assert import_globals_id_and_name("foo.bar", None) == ("foo", "foo")
+
+    def test_import_globals_id_and_name_for_import_as(self) -> None:
+        assert import_globals_id_and_name("foo.bar", "baz") == ("foo.bar", "baz")
+
+    def test_split_import_group_to_python_and_native_preserves_runs(self) -> None:
+        builder = make_builder(
+            native_modules={"pkg.alpha", "pkg.beta", "pkg.gamma"},
+            same_group_modules={"pkg.alpha", "pkg.beta", "pkg.gamma"},
+        )
+        group = [
+            Import([("pkg.alpha", None), ("py_mod", None)]),
+            Import([("pkg.beta", "beta_alias"), ("foreign.mod", None), ("pkg.gamma", None)]),
+        ]
+        group[0].line = 10
+        group[1].line = 20
+
+        result = split_import_group_to_python_and_native(builder, group)
+
+        assert result == [
+            ([("pkg.alpha", None, 10)], True),
+            ([("py_mod", None, 10)], False),
+            ([("pkg.beta", "beta_alias", 20)], True),
+            ([("foreign.mod", None, 20)], False),
+            ([("pkg.gamma", None, 20)], True),
+        ]
+
+    def test_group_consecutive_groups_by_kind_and_preserves_aliases(self) -> None:
+        buckets = group_consecutive(
+            [
+                (IMPORT_NATIVE_SUBMODULE, "a", "a"),
+                (IMPORT_NATIVE_SUBMODULE, "b", "b_alias"),
+                (IMPORT_NON_NATIVE, "c", "c"),
+                (IMPORT_NATIVE_ATTR, "d", "d_alias"),
+                (IMPORT_NATIVE_ATTR, "e", "e"),
+            ]
+        )
+
+        assert [(bucket.kind, bucket.names, bucket.as_names) for bucket in buckets] == [
+            (IMPORT_NATIVE_SUBMODULE, ["a", "b"], ["a", "b_alias"]),
+            (IMPORT_NON_NATIVE, ["c"], ["c"]),
+            (IMPORT_NATIVE_ATTR, ["d", "e"], ["d_alias", "e"]),
+        ]
+
+    def test_classify_import_from_groups_consecutive_kinds(self) -> None:
+        builder = make_builder(
+            native_modules={"pkg.native_a", "pkg.native_b"},
+            same_group_modules={"pkg.native_a", "pkg.native_b"},
+            graph={"pkg.native_a", "pkg.native_b", "pkg.foreign_a", "pkg.foreign_b"},
+        )
+
+        buckets = classify_import_from(
+            builder,
+            "pkg",
+            ["native_a", "native_b", "foreign_a", "foreign_b"],
+            ["native_a", "native_b_alias", "foreign_a", "foreign_b_alias"],
+            parent_is_native=True,
+        )
+
+        assert [(bucket.kind, bucket.names, bucket.as_names) for bucket in buckets] == [
+            (
+                IMPORT_NATIVE_SUBMODULE,
+                ["native_a", "native_b"],
+                ["native_a", "native_b_alias"],
+            ),
+            (
+                IMPORT_NON_NATIVE,
+                ["foreign_a", "foreign_b"],
+                ["foreign_a", "foreign_b_alias"],
+            ),
+        ]
+
+    def test_classify_import_from_treats_missing_name_under_native_parent_as_attr(self) -> None:
+        builder = make_builder(graph={"pkg.foreign"})
+
+        buckets = classify_import_from(
+            builder,
+            "pkg",
+            ["attr_name", "foreign"],
+            ["attr_alias", "foreign_alias"],
+            parent_is_native=True,
+        )
+
+        assert [(bucket.kind, bucket.names, bucket.as_names) for bucket in buckets] == [
+            (IMPORT_NATIVE_ATTR, ["attr_name"], ["attr_alias"]),
+            (IMPORT_NON_NATIVE, ["foreign"], ["foreign_alias"]),
+        ]
+
+    def test_classify_import_from_without_native_parent_never_uses_native_attr(self) -> None:
+        builder = make_builder()
+
+        buckets = classify_import_from(
+            builder,
+            "pkg",
+            ["attr_name"],
+            ["attr_alias"],
+            parent_is_native=False,
+        )
+
+        assert [(bucket.kind, bucket.names, bucket.as_names) for bucket in buckets] == [
+            (IMPORT_NON_NATIVE, ["attr_name"], ["attr_alias"])
+        ]