Add relative import support to static python visitors

Summary:
This is relevant to a pytree benchmark we want to run, becasue some of the
dependencies in numpy use relative imports.

Since a *lot* of perf-sensitive code will use numpy in ML, data science, and other
compute-heavy scenarios, anything relevant to numpy seems like a pretty good
feature to have.

The test coverage is pretty exhaustive, it might be overkill

Reviewed By: DinoV

Differential Revision: D96385657

fbshipit-source-id: 78e39c19df4b010ae39655343847b6ac95ce977e

Author: stroxler

cinderx/PythonLib/cinderx/compiler/static/declaration_visitor.py

@@ -261,9 +261,13 @@ def visitImport(self, node: Import) -> None:
                 )
 
     def visitImportFrom(self, node: ImportFrom) -> None:
-        mod_name = node.module
-        if not mod_name or node.level:
-            raise NotImplementedError("relative imports aren't supported")
+        if node.level:
+            mod_name = self._resolve_relative_import(node)
+        else:
+            mod_name = node.module
+        if not mod_name:
+            self.syntax_error("empty module name in import", node)
+            return
         for name in node.names:
             child_name = name.asname or name.name
             self.module.declare_import(

cinderx/PythonLib/cinderx/compiler/static/type_binder.py

@@ -1755,9 +1755,13 @@ def visitImport(self, node: Import) -> None:
                 self.declare_local(declaration_name, typ)
 
     def visitImportFrom(self, node: ImportFrom) -> None:
-        mod_name = node.module
-        if node.level or not mod_name:
-            raise NotImplementedError("relative imports aren't supported")
+        if node.level:
+            mod_name = self._resolve_relative_import(node)
+        else:
+            mod_name = node.module
+        if not mod_name:
+            self.syntax_error("empty module name in import", node)
+            return
 
         if mod_name == "__static__":
             for alias in node.names:

cinderx/PythonLib/cinderx/compiler/static/visitor.py

@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from ast import AST
+from ast import AST, ImportFrom
 from contextlib import contextmanager, nullcontext
 from typing import ContextManager, Generator, Generic, TYPE_CHECKING, TypeVar
 
@@ -81,6 +81,35 @@ def error_context(self, node: AST | None) -> ContextManager[None]:
             return nullcontext()
         return self.error_sink.error_context(self.filename, node)
 
+    def _resolve_relative_import(self, node: ImportFrom) -> str:
+        """Resolve a relative import to an absolute module name.
+
+        For __init__ modules, level=1 means "this package" (no stripping).
+        For non-__init__ modules, level=1 means "parent package" (strip one).
+        Each additional level strips one more package component.
+        """
+        level = node.level
+        parts = self.module_name.split(".")
+        # __init__ modules ARE the package, so relative imports start from
+        # the package itself rather than its parent.
+        is_package = self.filename.endswith("__init__.py")
+        strip = level if not is_package else level - 1
+        if strip > len(parts):
+            self.syntax_error(
+                "attempted relative import beyond top-level package", node
+            )
+            return ""
+        # Copy with list() when strip == 0 to avoid mutating `parts` via append below.
+        base_parts = parts[: len(parts) - strip] if strip > 0 else list(parts)
+        if node.module:
+            base_parts.append(node.module)
+        result = ".".join(base_parts)
+        if not result:
+            self.syntax_error(
+                "attempted relative import beyond top-level package", node
+            )
+        return result
+
     @contextmanager
     def temporary_error_sink(self, sink: ErrorSink) -> Generator[None, None, None]:
         orig_sink = self.error_sink

cinderx/PythonLib/test_cinderx/test_compiler/test_static/decl_visitor.py

@@ -295,3 +295,101 @@ def import_module(self, name: str, optimize: int) -> ModuleTable:
             compiler.compile(
                 "a", "a.py", ast.parse(dedent(codestr)), codestr, optimize=1
             )
+
+    # --- Relative import tests ---
+
+    def test_relative_import_preserves_static_type(self) -> None:
+        """from .foo import C in pkg.mod should resolve C statically,
+        producing INVOKE_FUNCTION bytecode instead of a dynamic call."""
+        foo_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            from .foo import C
+
+            def g():
+                x = C()
+                return x.f()
+        """
+        comp = self.compiler(**{"pkg.foo": foo_code, "pkg.mod": mod_code})
+        f = self.find_code(comp.compile_module("pkg.mod"), "g")
+        self.assertInBytecode(f, "INVOKE_FUNCTION", ((("pkg.foo", "C"), "f"), 1))
+
+    def test_relative_import_multi_level_preserves_static_type(self) -> None:
+        """from ..foo import C in pkg.sub.mod should resolve to pkg.foo.C."""
+        foo_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            from ..foo import C
+
+            def g():
+                x = C()
+                return x.f()
+        """
+        comp = self.compiler(**{"pkg.foo": foo_code, "pkg.sub.mod": mod_code})
+        f = self.find_code(comp.compile_module("pkg.sub.mod"), "g")
+        self.assertInBytecode(f, "INVOKE_FUNCTION", ((("pkg.foo", "C"), "f"), 1))
+
+    def test_relative_import_from_dot_no_module(self) -> None:
+        """from . import foo in pkg.mod should resolve foo as a module,
+        and chained attribute access foo.f() should produce INVOKE_FUNCTION."""
+        pkg_code = """
+            pass
+        """
+        foo_code = """
+            def f(x: int) -> int:
+                return x
+        """
+        mod_code = """
+            from . import foo
+
+            def g():
+                return foo.f(1)
+        """
+        comp = self.compiler(
+            **{"pkg": pkg_code, "pkg.foo": foo_code, "pkg.mod": mod_code}
+        )
+        f = self.find_code(comp.compile_module("pkg.mod"), "g")
+        self.assertInBytecode(f, "INVOKE_FUNCTION", ((("pkg.foo",), "f"), 1))
+
+    def test_relative_import_in_function_scope(self) -> None:
+        """Function-scope relative import should work via type_binder path."""
+        foo_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            def g():
+                from .foo import C
+                x = C()
+                return x.f()
+        """
+        comp = self.compiler(**{"pkg.foo": foo_code, "pkg.mod": mod_code})
+        f = self.find_code(comp.compile_module("pkg.mod"), "g")
+        self.assertInBytecode(f, "INVOKE_FUNCTION", ((("pkg.foo", "C"), "f"), 1))
+
+    def test_relative_import_type_checking(self) -> None:
+        """Relative import under TYPE_CHECKING should resolve for type annotations."""
+        foo_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            from typing import TYPE_CHECKING
+
+            if TYPE_CHECKING:
+                from .foo import C
+
+            def g(x: C):
+                return x.f()
+        """
+        comp = self.compiler(**{"pkg.foo": foo_code, "pkg.mod": mod_code})
+        f = self.find_code(comp.compile_module("pkg.mod"), "g")
+        self.assertInBytecode(f, "INVOKE_METHOD", ((("pkg.foo", "C"), "f"), 0))

cinderx/PythonLib/test_cinderx/test_compiler/test_static/imports.py

@@ -2,12 +2,108 @@
 
 # pyre-strict
 
+import ast
+import unittest
+
 from cinderx.compiler.pycodegen import PythonCodeGenerator
 from cinderx.compiler.static.types import TypedSyntaxError
+from cinderx.compiler.static.visitor import GenericVisitor
 
 from .common import StaticTestBase
 
 
+def _make_import_from(code: str) -> ast.ImportFrom:
+    """Parse a single 'from ... import ...' statement into an ImportFrom node."""
+    mod = ast.parse(code)
+    node = mod.body[0]
+    assert isinstance(node, ast.ImportFrom)
+    return node
+
+
+class ResolveRelativeImportTests(unittest.TestCase):
+    """Unit tests for GenericVisitor._resolve_relative_import name resolution.
+
+    These test the name manipulation logic directly, without needing the full
+    static compiler infrastructure. This lets us test __init__.py handling
+    which the TestCompiler infra can't exercise (its _get_filename never
+    produces __init__.py filenames).
+    """
+
+    def _resolve(self, module_name: str, filename: str, import_code: str) -> str:
+        """Call _resolve_relative_import with the given module context."""
+        node = _make_import_from(import_code)
+        # Create a minimal mock that has the attributes _resolve_relative_import needs:
+        # module_name, filename, and syntax_error.
+        visitor = object.__new__(GenericVisitor)
+        visitor.module_name = module_name
+        visitor.filename = filename
+        errors: list[str] = []
+        # pyre-ignore[8]: Attribute has type `(self: GenericVisitor[TVisitRet], msg: str, node: AST) -> None`
+        visitor.syntax_error = lambda msg, node: errors.append(msg)
+        result = visitor._resolve_relative_import(node)
+        if errors:
+            raise ValueError(errors[0])
+        return result
+
+    # --- Non-__init__ modules (regular .py files) ---
+
+    def test_single_dot_from_submodule(self) -> None:
+        # from .foo import C in pkg.mod -> pkg.foo
+        result = self._resolve("pkg.mod", "mod.py", "from .foo import C")
+        self.assertEqual(result, "pkg.foo")
+
+    def test_single_dot_no_module(self) -> None:
+        # from . import foo in pkg.mod -> base module is "pkg",
+        # then "foo" is resolved as a child of pkg by visitImportFrom.
+        result = self._resolve("pkg.mod", "mod.py", "from . import foo")
+        self.assertEqual(result, "pkg")
+
+    def test_double_dot(self) -> None:
+        # from ..foo import C in pkg.sub.mod -> pkg.foo
+        result = self._resolve("pkg.sub.mod", "mod.py", "from ..foo import C")
+        self.assertEqual(result, "pkg.foo")
+
+    def test_double_dot_no_module(self) -> None:
+        # from .. import foo in pkg.sub.mod -> base module is "pkg",
+        # then "foo" is resolved as a child of pkg by visitImportFrom.
+        result = self._resolve("pkg.sub.mod", "mod.py", "from .. import foo")
+        self.assertEqual(result, "pkg")
+
+    def test_beyond_top_level(self) -> None:
+        # from ...foo import C in pkg.mod -> error (only 2 parts, 3 dots)
+        with self.assertRaisesRegex(ValueError, "beyond top-level"):
+            self._resolve("pkg.mod", "mod.py", "from ...foo import C")
+
+    # --- __init__.py modules (package __init__ files) ---
+
+    def test_init_single_dot(self) -> None:
+        # from . import foo in pkg/__init__.py (module_name="pkg") -> base is "pkg",
+        # then "foo" is resolved as a child of pkg by visitImportFrom.
+        result = self._resolve("pkg", "__init__.py", "from . import foo")
+        self.assertEqual(result, "pkg")
+
+    def test_init_single_dot_with_module(self) -> None:
+        # from .sub import C in pkg/__init__.py (module_name="pkg") -> pkg.sub
+        result = self._resolve("pkg", "__init__.py", "from .sub import C")
+        self.assertEqual(result, "pkg.sub")
+
+    def test_init_double_dot(self) -> None:
+        # from .. import foo in pkg.sub/__init__.py (module_name="pkg.sub") -> base is "pkg",
+        # then "foo" is resolved as a child of pkg by visitImportFrom.
+        result = self._resolve("pkg.sub", "__init__.py", "from .. import foo")
+        self.assertEqual(result, "pkg")
+
+    def test_init_double_dot_with_module(self) -> None:
+        # from ..other import C in pkg.sub/__init__.py -> pkg.other
+        result = self._resolve("pkg.sub", "__init__.py", "from ..other import C")
+        self.assertEqual(result, "pkg.other")
+
+    def test_init_beyond_top_level(self) -> None:
+        # from .. import foo in top-level pkg/__init__.py -> error
+        with self.assertRaisesRegex(ValueError, "beyond top-level"):
+            self._resolve("pkg", "__init__.py", "from .. import foo")
+
+
 class ImportTests(StaticTestBase):
     def test_unknown_import_with_fallback_is_not_allowed(self) -> None:
         codestr = """

cinderx/PythonLib/test_cinderx/test_compiler/test_static/module.py

@@ -3,7 +3,7 @@
 # pyre-strict
 
 from cinderx.compiler.static.module_table import ModuleTableException
-from cinderx.compiler.static.types import ModuleInstance
+from cinderx.compiler.static.types import Class, ModuleInstance, TypedSyntaxError
 
 from .common import get_child, StaticTestBase, TestCompiler
 
@@ -314,6 +314,96 @@ class A:
         compiler = self.decl_visit(**{"a": acode, "b": bcode})
         compiler.compile_module("a")
 
+    # --- Relative import resolution tests ---
+
+    def test_relative_import_from_dot_module(self) -> None:
+        """from .foo import X in pkg.mod resolves to pkg.foo.X"""
+        foo_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            from .foo import C
+        """
+        compiler = self.decl_visit(**{"pkg.foo": foo_code, "pkg.mod": mod_code})
+
+        c = get_child(compiler.modules["pkg.mod"], "C")
+        self.assertIsNotNone(c)
+        assert isinstance(c, Class)
+        self.assertEqual(c.type_name.module, "pkg.foo")
+        self.assertEqual(c.type_name.qualname, "C")
+
+    def test_relative_import_from_dot_no_module(self) -> None:
+        """from . import foo in pkg.mod resolves to pkg.foo"""
+        foo_code = """
+            def f(x: int) -> int:
+                return x
+        """
+        mod_code = """
+            from . import foo
+        """
+        pkg_code = """
+            pass
+        """
+        compiler = self.decl_visit(
+            **{"pkg": pkg_code, "pkg.foo": foo_code, "pkg.mod": mod_code}
+        )
+
+        foo = get_child(compiler.modules["pkg.mod"], "foo")
+        self.assertIsNotNone(foo)
+        assert isinstance(foo, ModuleInstance)
+        self.assertEqual(foo.module_name, "pkg.foo")
+
+    def test_relative_import_multi_level(self) -> None:
+        """from .. import foo in pkg.sub.mod resolves to pkg.foo"""
+        foo_code = """
+            def f(x: int) -> int:
+                return x
+        """
+        mod_code = """
+            from .. import foo
+        """
+        pkg_code = """
+            pass
+        """
+        compiler = self.decl_visit(
+            **{"pkg": pkg_code, "pkg.foo": foo_code, "pkg.sub.mod": mod_code}
+        )
+
+        foo = get_child(compiler.modules["pkg.sub.mod"], "foo")
+        self.assertIsNotNone(foo)
+        assert isinstance(foo, ModuleInstance)
+        self.assertEqual(foo.module_name, "pkg.foo")
+
+    def test_relative_import_multi_level_with_module(self) -> None:
+        """from ..other import C in pkg.sub.mod resolves to pkg.other.C"""
+        other_code = """
+            class C:
+                def f(self) -> int:
+                    return 42
+        """
+        mod_code = """
+            from ..other import C
+        """
+        compiler = self.decl_visit(**{"pkg.other": other_code, "pkg.sub.mod": mod_code})
+
+        c = get_child(compiler.modules["pkg.sub.mod"], "C")
+        self.assertIsNotNone(c)
+        assert isinstance(c, Class)
+        self.assertEqual(c.type_name.module, "pkg.other")
+        self.assertEqual(c.type_name.qualname, "C")
+
+    def test_relative_import_beyond_top_level(self) -> None:
+        """Relative import with too many dots should error."""
+        mod_code = """
+            from ...foo import C
+        """
+        with self.assertRaisesRegex(
+            TypedSyntaxError, "attempted relative import beyond top-level package"
+        ):
+            self.decl_visit(**{"pkg.mod": mod_code})
+
     def test_actual_cyclic_reference(self) -> None:
         acode = """
             from b import B