jaclang 0.8.7__py3-none-any.whl → 0.8.8__py3-none-any.whl

jaclang/compiler/passes/main/tests/test_checker_pass.py

@@ -28,6 +28,18 @@ class TypeCheckerPassTests(TestCase):
                  ^^^^^^^^^^^^^^^^^^^^^^
         """, program.errors_had[1].pretty_print())
 
+    def test_float_types(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker_float.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            f: float = pi; # <-- OK
+            s: str = pi;   # <-- Error
+            ^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
+
     def test_infer_type_of_assignment(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("infer_type_assignment.jac"))
@@ -49,6 +61,17 @@ class TypeCheckerPassTests(TestCase):
           ^^^^^^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_imported_sym(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker/import_sym_test.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+          a: str = foo();  # <-- Ok
+          b: int = foo();  # <-- Error
+          ^^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_member_access_type_infered(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("member_access_type_inferred.jac"))
@@ -59,6 +82,22 @@ class TypeCheckerPassTests(TestCase):
           ^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_inherited_symbol(self) -> None:
+        program = JacProgram()
+        mod = program.compile(self.fixture_abs_path("checker_sym_inherit.jac"))
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 2)
+        self._assert_error_pretty_found("""
+          c.val = 42;     # <-- Ok
+          c.val = "str";  # <-- Error
+          ^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+        self._assert_error_pretty_found("""
+          l.name = "Simba";  # <-- Ok
+          l.name = 42;       # <-- Error
+          ^^^^^^^^^^^
+        """, program.errors_had[1].pretty_print())
+
     def test_import_symbol_type_infer(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("import_symbol_type_infer.jac"))
@@ -104,6 +143,145 @@ class TypeCheckerPassTests(TestCase):
             ^^^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_arity(self) -> None:
+        path = self.fixture_abs_path("checker_arity.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 3)
+        self._assert_error_pretty_found("""
+            f.first_is_self(f); # <-- Error
+                            ^
+        """, program.errors_had[0].pretty_print())
+        self._assert_error_pretty_found("""
+            f.with_default_args(1, 2, 3); # <-- Error
+                                      ^
+        """, program.errors_had[1].pretty_print())
+        self._assert_error_pretty_found("""
+            f.with_default_args();        # <-- Error
+            ^^^^^^^^^^^^^^^^^^^^^
+        """, program.errors_had[2].pretty_print())
+
+    def test_param_types(self) -> None:
+        path = self.fixture_abs_path("checker_param_types.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            foo(A()); # <-- Ok
+            foo(B()); # <-- Error
+                ^^^
+        """, program.errors_had[0].pretty_print())
+
+    def test_param_arg_match(self) -> None:
+        path = self.fixture_abs_path("checker_param_types.jac")
+        program = JacProgram()
+        path = self.fixture_abs_path("checker_arg_param_match.jac")
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 13)
+
+        expected_errors = [
+            """
+            Not all required parameters were provided in the function call: 'a'
+                     f = Foo();
+                     f.bar();
+                     ^^^^^^^
+            """,
+            """
+            Too many positional arguments
+                     f.bar();
+                     f.bar(1);
+                     f.bar(1, 2);
+                              ^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'self', 'a'
+                     f.bar(1, 2);
+                     f.baz();
+                     ^^^^^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'a'
+                     f.baz();
+                     f.baz(1);
+                     ^^^^^^^^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'f'
+                     foo(1, 2, d=3, e=4, f=5, c=4);  # order does not matter for named
+                     foo(1, 2, 3, d=4, e=5, g=7, h=8); # missing argument 'f'
+                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+            """,
+            """
+            Positional only parameter 'b' cannot be matched with a named argument
+                     foo(1, 2, 3, d=4, e=5, g=7, h=8); # missing argument 'f'
+                     foo(1, b=2, c=3, d=4, e=5, f=6); # b is positional only
+                            ^^^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                           ^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                              ^
+            """,
+            """
+            Too many positional arguments
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                                                 ^
+            """,
+            """
+            Parameter 'c' already matched
+                     bar(1, 2, 3, 4, 5, f=6);
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                     bar(1, 2, 3, 4, 5, 6, c=3); # already matched
+                                           ^^^
+            """,
+            """
+            Named argument 'h' does not match any parameter
+                     bar(1, 2, 3, 4, 5, 6, 7, 8, 9); # too many args
+                     bar(1, 2, 3, 4, 5, 6, c=3); # already matched
+                     bar(1, 2, 3, 4, 5, 6, h=1); # h is not matched
+                                           ^^^
+            """,
+            """
+            Too many positional arguments
+                     baz(a=1, b=2);
+                     baz(1, b=2); # a can be both positional and keyword
+                     baz(1, 2);  # 'b' can only be keyword arg
+                            ^
+            """,
+            """
+            Not all required parameters were provided in the function call: 'b'
+                     baz(a=1, b=2);
+                     baz(1, b=2); # a can be both positional and keyword
+                     baz(1, 2);  # 'b' can only be keyword arg
+                     ^^^^^^^^^
+            """,
+        ]
+
+        for i, expected in enumerate(expected_errors):
+            self._assert_error_pretty_found(expected, program.errors_had[i].pretty_print())
+
+    def test_self_type_inference(self) -> None:
+        path = self.fixture_abs_path("checker_self_type.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+          x: str = self.i; # <-- Error
+          ^^^^^^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_binary_op(self) -> None:
         program = JacProgram()
         mod = program.compile(self.fixture_abs_path("checker_binary_op.jac"))
@@ -140,6 +318,19 @@ class TypeCheckerPassTests(TestCase):
             ^^^^^^^^^^^^^^
         """, program.errors_had[0].pretty_print())
 
+    def test_checker_cat_is_animal(self) -> None:
+        path = self.fixture_abs_path("checker_cat_is_animal.jac")
+        program = JacProgram()
+        mod = program.compile(path)
+        TypeCheckPass(ir_in=mod, prog=program)
+        self.assertEqual(len(program.errors_had), 1)
+        self._assert_error_pretty_found("""
+            animal_func(cat);        # <-- Ok
+            animal_func(lion);       # <-- Ok
+            animal_func(not_animal); # <-- Error
+                        ^^^^^^^^^^
+        """, program.errors_had[0].pretty_print())
+
     def test_checker_import_missing_module(self) -> None:
         path = self.fixture_abs_path("checker_import_missing_module.jac")
         program = JacProgram()

jaclang/compiler/passes/main/tests/test_predynamo_pass.py

@@ -0,0 +1,57 @@
+"""Test ast build pass module."""
+
+import io
+import os
+import sys
+
+from jaclang.compiler.program import JacProgram
+from jaclang.utils.test import TestCase
+from jaclang.settings import settings
+from jaclang.compiler.passes.main import PreDynamoPass
+
+class PreDynamoPassTests(TestCase):
+    """Test pass module."""
+
+    TargetPass = PreDynamoPass
+
+    def setUp(self) -> None:
+        """Set up test."""
+        return super().setUp()
+    
+    def test_predynamo_where_assign(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        os.environ["JAC_PREDYNAMO_PASS"] = "True"
+        settings.load_env_vars()
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_where_assign.jac"))
+        sys.stdout = sys.__stdout__
+        self.assertIn("torch.where", code_gen.unparse())
+        os.environ["JAC_PREDYNAMO_PASS"] = "false"
+        settings.load_env_vars()
+        
+    def test_predynamo_where_return(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        os.environ["JAC_PREDYNAMO_PASS"] = "True"
+        settings.load_env_vars()
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_where_return.jac"))
+        sys.stdout = sys.__stdout__
+        self.assertIn("torch.where", code_gen.unparse())
+        os.environ["JAC_PREDYNAMO_PASS"] = "false"
+        settings.load_env_vars()
+
+    def test_predynamo_fix3(self) -> None:
+        """Test torch.where transformation."""
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        os.environ["JAC_PREDYNAMO_PASS"] = "True"
+        settings.load_env_vars()
+        code_gen = JacProgram().compile(self.fixture_abs_path("predynamo_fix3.jac"))
+        sys.stdout = sys.__stdout__
+        unparsed_code = code_gen.unparse()
+        self.assertIn("__inv_freq = torch.where(", unparsed_code)
+        self.assertIn("self.register_buffer('inv_freq', __inv_freq, persistent=False);", unparsed_code)
+        os.environ["JAC_PREDYNAMO_PASS"] = "false"
+        settings.load_env_vars()

jaclang/compiler/passes/main/type_checker_pass.py

@@ -9,92 +9,34 @@ Reference:
     craizy_type_expr branch: type_checker_pass.py
 """
 
-import ast as py_ast
-import os
-
 import jaclang.compiler.unitree as uni
 from jaclang.compiler.passes import UniPass
-from jaclang.compiler.type_system.type_evaluator import TypeEvaluator
-from jaclang.runtimelib.utils import read_file_with_encoding
-
-from .pyast_load_pass import PyastBuildPass
-from .sym_tab_build_pass import SymTabBuildPass
+from jaclang.compiler.type_system import types as jtypes
 
 
 class TypeCheckPass(UniPass):
     """Type checker pass for JacLang."""
 
-    # NOTE: This is done in the binder pass of pyright, however I'm doing this
-    # here, cause this will be the entry point of the type checker and we're not
-    # relying on the binder pass at the moment and we can go back to binder pass
-    # in the future if we needed it.
-    _BUILTINS_STUB_FILE_PATH = os.path.join(
-        os.path.dirname(__file__),
-        "../../../vendor/typeshed/stdlib/builtins.pyi",
-    )
-
-    # Cache the builtins module once it parsed.
-    _BUILTINS_MODULE: uni.Module | None = None
-
-    # REVIEW: Making the evaluator a static (singleton) variable to make sure only one
-    # instance is used across mulitple compilation units. This can also be attached to an
-    # attribute of JacProgram, however the evaluator is a temproary object that we dont
-    # want bound to the program for long term, Also the program is the one that will be
-    # dumped in the compiled bundle.
-    _EVALUATOR: TypeEvaluator | None = None
-
     def before_pass(self) -> None:
         """Initialize the checker pass."""
-        self._load_builtins_stub_module()
+        self.evaluator = self.prog.get_type_evaluator()
+        self.evaluator.diagnostic_callback = self._add_diagnostic
         self._insert_builtin_symbols()
 
-    @property
-    def evaluator(self) -> TypeEvaluator:
-        """Return the type evaluator."""
-        if TypeCheckPass._EVALUATOR is None:
-            assert TypeCheckPass._BUILTINS_MODULE is not None
-            TypeCheckPass._EVALUATOR = TypeEvaluator(
-                builtins_module=TypeCheckPass._BUILTINS_MODULE,
-                program=self.prog,
-            )
-        return TypeCheckPass._EVALUATOR
+    def _add_diagnostic(self, node: uni.UniNode, message: str, warning: bool) -> None:
+        """Add a diagnostic message to the pass."""
+        if warning:
+            self.log_warning(message, node)
+        else:
+            self.log_error(message, node)
 
     # --------------------------------------------------------------------------
     # Internal helper functions
     # --------------------------------------------------------------------------
 
-    def _binding_builtins(self) -> bool:
-        """Return true if we're binding the builtins stub file."""
-        return self.ir_in == TypeCheckPass._BUILTINS_MODULE
-
-    def _load_builtins_stub_module(self) -> None:
-        """Return the builtins stub module.
-
-        This will parse and cache the stub file and return the cached module on
-        subsequent calls.
-        """
-        if self._binding_builtins() or TypeCheckPass._BUILTINS_MODULE is not None:
-            return
-
-        if not os.path.exists(TypeCheckPass._BUILTINS_STUB_FILE_PATH):
-            raise FileNotFoundError(
-                f"Builtins stub file not found at {TypeCheckPass._BUILTINS_STUB_FILE_PATH}"
-            )
-
-        file_content = read_file_with_encoding(TypeCheckPass._BUILTINS_STUB_FILE_PATH)
-        uni_source = uni.Source(file_content, TypeCheckPass._BUILTINS_STUB_FILE_PATH)
-        mod = PyastBuildPass(
-            ir_in=uni.PythonModuleAst(
-                py_ast.parse(file_content),
-                orig_src=uni_source,
-            ),
-            prog=self.prog,
-        ).ir_out
-        SymTabBuildPass(ir_in=mod, prog=self.prog)
-        TypeCheckPass._BUILTINS_MODULE = mod
-
     def _insert_builtin_symbols(self) -> None:
-        if self._binding_builtins():
+        # Don't insert builtin symbols into the builtin module itself.
+        if self.ir_in == self.evaluator.builtins_module:
             return
 
         # TODO: Insert these symbols.
@@ -104,20 +46,34 @@ class TypeCheckPass(UniPass):
         # '__name__', '__loader__', '__package__', '__spec__', '__path__',
         # '__file__', '__cached__', '__dict__', '__annotations__',
         # '__builtins__', '__doc__',
-        assert (
-            TypeCheckPass._BUILTINS_MODULE is not None
-        ), "Builtins module is not loaded"
         if self.ir_in.parent_scope is not None:
             self.log_info("Builtins module is already bound, skipping.")
             return
         # Review: If we ever assume a module cannot have a parent scope, this will
         # break that contract.
-        self.ir_in.parent_scope = TypeCheckPass._BUILTINS_MODULE
+        self.ir_in.parent_scope = self.evaluator.builtins_module
 
     # --------------------------------------------------------------------------
     # Ast walker hooks
     # --------------------------------------------------------------------------
 
+    def enter_ability(self, node: uni.Ability) -> None:
+        """Enter an ability node."""
+        # If the node has @staticmethod decorator, mark it as static method.
+        # this is needed since ast raised from python does not have this info.
+        for decor in node.decorators or []:
+            ty = self.evaluator.get_type_of_expression(decor)
+            if isinstance(ty, jtypes.ClassType) and ty.is_builtin("staticmethod"):
+                node.is_static = True
+                break
+
+    def exit_import(self, node: uni.Import) -> None:
+        """Exit an import node."""
+        if node.from_loc:
+            for item in node.items:
+                if isinstance(item, uni.ModuleItem):
+                    self.evaluator.get_type_of_module_item(item)
+
     def exit_assignment(self, node: uni.Assignment) -> None:
         """Pyright: Checker.visitAssignment(node: AssignmentNode): boolean."""
         # TODO: In pyright this logic is present at evaluateTypesForAssignmentStatement