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

jaclang/compiler/type_system/type_evaluator.py

@@ -6,20 +6,30 @@ PyrightReference:
     packages/pyright-internal/src/analyzer/typeEvaluatorTypes.ts
 """
 
+import ast as py_ast
+import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, cast
+from typing import Callable, TYPE_CHECKING, cast
 
 import jaclang.compiler.unitree as uni
+from jaclang.compiler import TOKEN_MAP
+from jaclang.compiler.constant import Tokens as Tok
+from jaclang.compiler.passes.main.pyast_load_pass import PyastBuildPass
+from jaclang.compiler.passes.main.sym_tab_build_pass import SymTabBuildPass
 from jaclang.compiler.type_system import types
+from jaclang.runtimelib.utils import read_file_with_encoding
 
 if TYPE_CHECKING:
     from jaclang.compiler.program import JacProgram
 
 from . import operations
-from .type_utils import ClassMember
+from . import type_utils
 from .types import TypeBase
 
+# The callback type definition for the diagnostic messages.
+DiagnosticCallback = Callable[[uni.UniNode, str, bool], None]
+
 
 @dataclass
 class PrefetchedTypes:
@@ -35,6 +45,7 @@ class PrefetchedTypes:
     tuple_class: TypeBase | None = None
     bool_class: TypeBase | None = None
     int_class: TypeBase | None = None
+    float_class: TypeBase | None = None
     str_class: TypeBase | None = None
     dict_class: TypeBase | None = None
     module_type_class: TypeBase | None = None
@@ -60,10 +71,35 @@ class SymbolResolutionStackEntry:
     partial_type: TypeBase | None = None
 
 
+@dataclass
+class MatchArgsToParamsResult:
+    """Result of matching arguments to parameters."""
+
+    # FIXME: This class implementation is modified from pyright to make it
+    # simple and work for now, however this needs to be revisited and
+    # implemented properly.
+    arg_params: dict[uni.Expr | uni.KWPair, types.Parameter | None]
+
+    overload: types.FunctionType | None = None
+    argument_errors: bool = False
+
+
 class TypeEvaluator:
     """Type evaluator for JacLang."""
 
-    def __init__(self, builtins_module: uni.Module, program: "JacProgram") -> None:
+    # 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",
+    )
+
+    def __init__(
+        self,
+        program: "JacProgram",
+    ) -> None:
         """Initialize the type evaluator with prefetched types.
 
         Implementation Note:
@@ -74,10 +110,68 @@ class TypeEvaluator:
         in some place then it will not be available in the evaluator, So we
         are prefetching the builtins at the constructor level once.
         """
-        self.symbol_resolution_stack: list[SymbolResolutionStackEntry] = []
-        self.builtins_module = builtins_module
         self.program = program
+        self.symbol_resolution_stack: list[SymbolResolutionStackEntry] = []
+        self.builtins_module = self._load_builtins_stub_module()
         self.prefetch = self._prefetch_types()
+        self.diagnostic_callback: DiagnosticCallback | None = None
+
+    def _load_builtins_stub_module(self) -> uni.Module:
+        """Load and return the builtins stub module."""
+        if not os.path.exists(TypeEvaluator._BUILTINS_STUB_FILE_PATH):
+            raise FileNotFoundError(
+                f"Builtins stub file not found at {TypeEvaluator._BUILTINS_STUB_FILE_PATH}"
+            )
+        file_content = read_file_with_encoding(TypeEvaluator._BUILTINS_STUB_FILE_PATH)
+        uni_source = uni.Source(file_content, TypeEvaluator._BUILTINS_STUB_FILE_PATH)
+        mod = PyastBuildPass(
+            ir_in=uni.PythonModuleAst(
+                py_ast.parse(file_content),
+                orig_src=uni_source,
+            ),
+            prog=self.program,
+        ).ir_out
+        SymTabBuildPass(ir_in=mod, prog=self.program)
+        return mod
+
+    def _get_builtin_type(self, name: str) -> TypeBase:
+        """Return the built-in type with the given name."""
+        if (symbol := self.builtins_module.lookup(name)) is not None:
+            return self.get_type_of_symbol(symbol)
+        return types.UnknownType()
+
+    def _prefetch_types(self) -> "PrefetchedTypes":
+        """Return the prefetched types for the type evaluator."""
+        return PrefetchedTypes(
+            # TODO: Pyright first try load NoneType from typeshed and if it cannot
+            # then it set to unknown type.
+            none_type_class=types.UnknownType(),
+            object_class=self._get_builtin_type("object"),
+            type_class=self._get_builtin_type("type"),
+            # union_type_class=
+            # awaitable_class=
+            # function_class=
+            # method_class=
+            tuple_class=self._get_builtin_type("tuple"),
+            bool_class=self._get_builtin_type("bool"),
+            int_class=self._get_builtin_type("int"),
+            float_class=self._get_builtin_type("float"),
+            str_class=self._get_builtin_type("str"),
+            dict_class=self._get_builtin_type("dict"),
+            # module_type_class=
+            # typed_dict_class=
+            # typed_dict_private_class=
+            # supports_keys_and_get_item_class=
+            # mapping_class=
+            # template_class=
+        )
+
+    def add_diagnostic(
+        self, node: uni.UniNode, message: str, warning: bool = False
+    ) -> None:
+        """Add a diagnostic message to the program."""
+        if self.diagnostic_callback:
+            self.diagnostic_callback(node, message, warning)
 
     # -------------------------------------------------------------------------
     # Symbol resolution stack
@@ -229,15 +323,25 @@ class TypeEvaluator:
         if node.name_spec.type is not None:
             return cast(types.ClassType, node.name_spec.type)
 
+        base_classes: list[TypeBase] = []
+        for base_class in node.base_classes or []:
+            base_class_type = self.get_type_of_expression(base_class)
+            base_classes.append(base_class_type)
+        is_builtin_class = node.find_parent_of_type(uni.Module) == self.builtins_module
+
         cls_type = types.ClassType(
             types.ClassType.ClassDetailsShared(
                 class_name=node.name_spec.sym_name,
                 symbol_table=node,
-                # TODO: Resolve the base class expression and pass them here.
+                base_classes=base_classes,
+                is_builtin_class=is_builtin_class,
             ),
             flags=types.TypeFlags.Instantiable,
         )
 
+        # Compute the MRO for the class.
+        type_utils.compute_mro_linearization(cls_type)
+
         # Cache the type, pyright is doing invalidateTypeCacheIfCanceled()
         # we're not doing that any time sooner.
         node.name_spec.type = cls_type
@@ -252,17 +356,61 @@ class TypeEvaluator:
             node.name_spec.type = types.UnknownType()
             return node.name_spec.type
 
-        if not isinstance(node.signature.return_type, uni.Expr):
-            node.name_spec.type = types.UnknownType()
-            return node.name_spec.type
+        return_type: TypeBase
+        if isinstance(node.signature.return_type, uni.Expr):
+            return_type = self._convert_to_instance(
+                self.get_type_of_expression(node.signature.return_type)
+            )
+        else:
+            return_type = types.UnknownType()
+
+        # Define helper function for parameter conversion.
+        def _get_param_category(param: uni.ParamVar) -> types.ParameterCategory:
+            if param.is_vararg:
+                return types.ParameterCategory.ArgsList
+            if param.is_kwargs:
+                return types.ParameterCategory.KwargsDict
+            return types.ParameterCategory.Positional
+
+        # Define helper function for parameter kind conversion.
+        def _convert_param_kind(kind: uni.ParamKind) -> types.ParamKind:
+            match kind:
+                case uni.ParamKind.POSONLY:
+                    return types.ParamKind.POSONLY
+                case uni.ParamKind.NORMAL:
+                    return types.ParamKind.NORMAL
+                case uni.ParamKind.VARARG:
+                    return types.ParamKind.VARARG
+                case uni.ParamKind.KWONLY:
+                    return types.ParamKind.KWONLY
+                case uni.ParamKind.KWARG:
+                    return types.ParamKind.KWARG
+            return types.ParamKind.NORMAL
+
+        parameters: list[types.Parameter] = []
+        for idx, param in enumerate(node.signature.get_parameters()):
+            # TODO: Set parameter category for *args, and **kwargs
+            param_type: TypeBase | None = None
+
+            if param.type_tag:
+                param_type_cls = self.get_type_of_expression(param.type_tag.tag)
+                param_type = self._convert_to_instance(param_type_cls)
+
+            parameters.append(
+                types.Parameter(
+                    name=param.name.value,
+                    category=_get_param_category(param),
+                    param_type=param_type,
+                    default_value=param.value,
+                    is_self=(idx == 0 and self._is_expr_self(param.name)),
+                    param_kind=_convert_param_kind(param.param_kind),
+                )
+            )
 
-        return_type = self._convert_to_instance(
-            self.get_type_of_expression(node.signature.return_type)
-        )
         func_type = types.FunctionType(
             func_name=node.name_spec.sym_name,
             return_type=return_type,
-            parameters=[],  # TODO:
+            parameters=parameters,
         )
 
         node.name_spec.type = func_type
@@ -285,6 +433,11 @@ class TypeEvaluator:
         assert self.prefetch.int_class is not None
         return self.prefetch.int_class
 
+    def get_type_of_float(self, node: uni.Float) -> TypeBase:
+        """Return the effective type of the float."""
+        assert self.prefetch.float_class is not None
+        return self.prefetch.float_class
+
     # Pyright equivalent function name = getTypeOfExpression();
     def get_type_of_expression(self, node: uni.Expr) -> TypeBase:
         """Return the effective type of the expression."""
@@ -311,9 +464,6 @@ class TypeEvaluator:
             # NOTE: For now if we don't have the type info, we assume it's compatible.
             # For strict mode we should disallow usage of unknown unless explicitly ignored.
             return True
-        # FIXME: This logic is not valid, just here as a stub.
-        if types.TypeCategory.Unknown in (src_type.category, dest_type.category):
-            return True
 
         if src_type == dest_type:
             return True
@@ -323,8 +473,7 @@ class TypeEvaluator:
             assert isinstance(src_type, types.ClassType)
             return self._assign_class(src_type, dest_type)
 
-        # FIXME: This is temporary.
-        return src_type == dest_type
+        return False
 
     # TODO: This should take an argument list as parameter.
     def get_type_of_magic_method_call(
@@ -356,39 +505,23 @@ class TypeEvaluator:
         if src_type.shared == dest_type.shared:
             return True
 
-        # TODO: Search base classes and everything else pyright is doing.
-        return False
+        # Check if src class is a subclass of dest class.
+        for base_cls in src_type.shared.mro:
+            if base_cls.shared == dest_type.shared:
+                return True
 
-    def _prefetch_types(self) -> "PrefetchedTypes":
-        """Return the prefetched types for the type evaluator."""
-        return PrefetchedTypes(
-            # TODO: Pyright first try load NoneType from typeshed and if it cannot
-            # then it set to unknown type.
-            none_type_class=types.UnknownType(),
-            object_class=self._get_builtin_type("object"),
-            type_class=self._get_builtin_type("type"),
-            # union_type_class=
-            # awaitable_class=
-            # function_class=
-            # method_class=
-            tuple_class=self._get_builtin_type("tuple"),
-            bool_class=self._get_builtin_type("bool"),
-            int_class=self._get_builtin_type("int"),
-            str_class=self._get_builtin_type("str"),
-            dict_class=self._get_builtin_type("dict"),
-            # module_type_class=
-            # typed_dict_class=
-            # typed_dict_private_class=
-            # supports_keys_and_get_item_class=
-            # mapping_class=
-            # template_class=
-        )
+        # Everything is assignable to an object.
+        if dest_type.is_builtin("object"):
+            # TODO: Invariance not handled yet
+            # invariant contexts to avoid list[int] <: list[object] errors.
+            return True
 
-    def _get_builtin_type(self, name: str) -> TypeBase:
-        """Return the built-in type with the given name."""
-        if (symbol := self.builtins_module.lookup(name)) is not None:
-            return self.get_type_of_symbol(symbol)
-        return types.UnknownType()
+        # Integers can be used where floats are expected.
+        if src_type.is_builtin("int") and dest_type.is_builtin("float"):
+            return True
+
+        # TODO: Search base classes and everything else pyright is doing.
+        return False
 
     # This function is a combination of the bellow pyright functions.
     #  - getDeclaredTypeOfSymbol
@@ -414,6 +547,11 @@ class TypeEvaluator:
             case uni.Ability():
                 return self.get_type_of_ability(node)
 
+            case uni.ParamVar():
+                if node.type_tag:
+                    annotation_type = self.get_type_of_expression(node.type_tag.tag)
+                    return self._convert_to_instance(annotation_type)
+
             # This actually defined in the function getTypeForDeclaration();
             # Pyright has DeclarationType.Variable.
             case uni.Name():
@@ -451,6 +589,9 @@ class TypeEvaluator:
             case uni.Int():
                 return self._convert_to_instance(self.get_type_of_int(expr))
 
+            case uni.Float():
+                return self._convert_to_instance(self.get_type_of_float(expr))
+
             case uni.AtomTrailer():
                 # NOTE: Pyright is using CFG to figure out the member type by narrowing the base
                 # type and filtering the members. We're not doing that anytime sooner.
@@ -496,25 +637,18 @@ class TypeEvaluator:
                 return self.get_type_of_expression(expr.value)
 
             case uni.FuncCall():
-                caller_type = self.get_type_of_expression(expr.target)
-                if isinstance(caller_type, types.FunctionType):
-                    return caller_type.return_type or types.UnknownType()
-                if (
-                    isinstance(caller_type, types.ClassType)
-                    and caller_type.is_instantiable_class()
-                ):
-                    return caller_type.clone_as_instance()
-                if caller_type.is_class_instance():
-                    magic_call_ret = self.get_type_of_magic_method_call(
-                        caller_type, "__call__"
-                    )
-                    if magic_call_ret:
-                        return magic_call_ret
+                return self.validate_call_args(expr)
 
             case uni.BinaryExpr():
                 return operations.get_type_of_binary_operation(self, expr)
 
             case uni.Name():
+                # NOTE: For self's type pyright is getting the first parameter of a method and
+                # the name can be anything not just self, however we don't have the first parameter
+                # and self is a keyword, we need to do it in this way.
+                if self._is_expr_self(expr):
+                    return self._get_type_of_self(expr)
+
                 if symbol := expr.sym_tab.lookup(expr.value, deep=True):
                     expr.sym = symbol
                     return self.get_type_of_symbol(symbol)
@@ -522,6 +656,48 @@ class TypeEvaluator:
             # TODO: More expressions.
         return types.UnknownType()
 
+    # -----------------------------------------------------------------------------
+    # Helper functions
+    # -----------------------------------------------------------------------------
+
+    def _is_expr_self(self, expr: uni.Expr) -> bool:
+        """Check if the expression is Name that is 'self' and in the method context."""
+        if (
+            isinstance(expr, uni.Name)
+            and (expr.value == TOKEN_MAP[Tok.KW_SELF])
+            and (fn := self._get_enclosing_method(expr))
+            and (not fn.is_static)
+        ):
+            return True
+        return False
+
+    def _get_enclosing_function(self, node: uni.UniNode) -> uni.Ability | None:
+        """Get the enclosing function (ability) of the given node."""
+        if (impl := node.find_parent_of_type(uni.ImplDef)) and (
+            isinstance(impl.decl_link, uni.Ability)
+        ):
+            return impl.decl_link
+        return node.find_parent_of_type(uni.Ability)
+
+    def _get_enclosing_method(self, node: uni.UniNode) -> uni.Ability | None:
+        """Get the enclosing method (ability) of the given node."""
+        enclosing_fn = self._get_enclosing_function(node)
+        while enclosing_fn and (not enclosing_fn.is_method):
+            enclosing_fn = self._get_enclosing_function(enclosing_fn)
+        if enclosing_fn and enclosing_fn.is_method:
+            return enclosing_fn
+        return None
+
+    def _get_type_of_self(self, node: uni.Name) -> TypeBase:
+        """Return the effective type of self."""
+        if method := self._get_enclosing_method(node):
+            cls = method.method_owner
+            if isinstance(cls, uni.Archetype):
+                return self.get_type_of_class(cls).clone_as_instance()
+            if isinstance(cls, uni.Enum):
+                pass  # TODO: Implement type from enum.
+        return types.UnknownType()
+
     def _convert_to_instance(self, jtype: TypeBase) -> TypeBase:
         """Convert a class type to an instance type."""
         # TODO: Grep pyright "Handle type[x] as a special case." They handle `type[x]` as a special case:
@@ -536,7 +712,7 @@ class TypeEvaluator:
 
     def _lookup_class_member(
         self, base_type: types.ClassType, member: str
-    ) -> ClassMember | None:
+    ) -> type_utils.ClassMember | None:
         """Lookup the class member type."""
         assert self.prefetch.int_class is not None
         # FIXME: Pyright's way: Implement class member iterator (based on mro and the multiple inheritance)
@@ -544,13 +720,14 @@ class TypeEvaluator:
 
         # NOTE: This is a simple implementation to make it work and more robust implementation will
         # be done in a future PR.
-        if sym := base_type.lookup_member_symbol(member):
-            return ClassMember(sym, base_type)
+        for cls in base_type.shared.mro:
+            if sym := cls.lookup_member_symbol(member):
+                return type_utils.ClassMember(sym, cls)
         return None
 
     def _lookup_object_member(
         self, base_type: types.ClassType, member: str
-    ) -> ClassMember | None:
+    ) -> type_utils.ClassMember | None:
         """Lookup the object member type."""
         assert self.prefetch.int_class is not None
         if base_type.is_class_instance():
@@ -558,3 +735,110 @@ class TypeEvaluator:
             # TODO: We need to implement Member lookup flags and set SkipInstanceMember to 0.
             return self._lookup_class_member(base_type, member)
         return None
+
+    def match_args_to_params(
+        self, expr: uni.FuncCall, func_type: types.FunctionType
+    ) -> MatchArgsToParamsResult:
+        """
+        Match arguments passed to a function to the corresponding parameters in that function.
+
+        This matching is done based on positions and keywords. Type evaluation and
+        validation is left to the caller.
+        This logic is based on PEP 3102: https://www.python.org/dev/peps/pep-3102/
+        """
+        arg_params: dict[uni.Expr | uni.KWPair, types.Parameter | None] = {}
+        argument_errors = False
+
+        params_to_match = func_type.parameters.copy()
+
+        # Skip `self` for method calls.
+        if len(func_type.parameters) >= 1 and func_type.parameters[0].is_self:
+            params_to_match.pop(0)
+
+        # Create a tracker for parameter assignment.
+        param_tracker = type_utils.ParamAssignmentTracker(params_to_match)
+
+        # We iterate over the arguments and match with the parameter, the param_tracker will
+        # keep track of the matched parameters and unmatched required parameters.
+        #
+        # Tracker: p1, p2, /, p3, p4,   *args,       | p6,   **kwargs
+        #          ^   ^      ^    ^    ^^   ^       | ^       ^^
+        #          |   |      |    |    | \__ \__    | |       | \________
+        # Args:    a1, a2,   a3, p4=a4, a5, a6, *a7, | p6=a8, p_kw1=a9, p_kw2=a10
+        #         '--------------------------------' | '------------------------'
+        #          We match positional with          |  We match named arguments with
+        #          tracked parameter index.          |  param name lookup.
+        #
+        for arg in expr.params:
+            try:
+                if isinstance(arg, uni.KWPair):
+                    # Match parameter based on name lookup.
+                    matching_param = param_tracker.match_named_argument(arg)
+                    arg_params[arg] = matching_param
+                else:  # Match parameter based on the position of the argument.
+                    matching_param = param_tracker.match_positional_argument(arg)
+                    arg_params[arg] = matching_param
+            except Exception as e:
+                self.add_diagnostic(arg, str(e))
+                argument_errors = True
+
+        if unmatched_params := param_tracker.get_unmatched_required_params():
+            names = ", ".join(f"'{p.name}'" for p in unmatched_params)
+            argument_errors = True
+            self.add_diagnostic(
+                expr,
+                f"Not all required parameters were provided in the function call: {names}",
+            )
+
+        return MatchArgsToParamsResult(
+            arg_params=arg_params, argument_errors=argument_errors
+        )
+
+    def validate_call_args(self, expr: uni.FuncCall) -> TypeBase:
+        """
+        Validate that the arguments can be assigned to the call's parameter list.
+
+        Specializes the call based on arg types, and returns the specialized
+        type of the return value. If it detects an error along the way, it emits
+        a diagnostic and sets argumentErrors to true.
+        """
+        caller_type = self.get_type_of_expression(expr.target)
+        if isinstance(caller_type, types.FunctionType):
+            arg_param_match = self.match_args_to_params(expr, caller_type)
+            if not arg_param_match.argument_errors:
+                self.validate_arg_types(arg_param_match)
+            return caller_type.return_type or types.UnknownType()
+
+        if (
+            isinstance(caller_type, types.ClassType)
+            and caller_type.is_instantiable_class()
+        ):
+            # TODO: validate args for __init__()
+            return caller_type.clone_as_instance()
+
+        if caller_type.is_class_instance():
+            # TODO: validate args.
+            magic_call_ret = self.get_type_of_magic_method_call(caller_type, "__call__")
+            if magic_call_ret:
+                return magic_call_ret
+
+        return types.UnknownType()
+
+    def validate_arg_types(
+        self,
+        args: MatchArgsToParamsResult,
+    ) -> None:
+        """Validate that the argument types can be assigned to the parameter types."""
+        for arg, param in args.arg_params.items():
+            if param is None or param.param_type is None:
+                continue
+            if isinstance(arg, uni.KWPair):
+                arg_type = self.get_type_of_expression(arg.value)
+            else:
+                arg_type = self.get_type_of_expression(arg)
+
+            if not self.assign_type(arg_type, param.param_type):
+                self.add_diagnostic(
+                    arg,
+                    f"Cannot assign {arg_type} to parameter '{param.name}' of type {param.param_type}",
+                )