feat(checker): resolve overloads with subtypes

try to find the most specific overload if multiple matches are found

examples/01_simple_type_checking/05_functions.py

@@ -0,0 +1,28 @@
+def incr(value: int):
+    return value + 1
+
+
+def decr(value: int):
+    return value - 1
+
+
+def foo(a: int, /, b: float, *, c: str):
+    return True
+
+
+r1 = foo()  # foo() missing 2 required positional arguments: 'a' and 'b'
+r2 = foo(1)  # foo() missing 1 required positional argument: 'b'
+r3 = foo(1, 2.0)  # foo() missing 1 required keyword-only argument: 'c'
+r4 = foo(1, b=2.0)  # foo() missing 1 required keyword-only argument: 'c'
+r5 = foo(1, 2.0, "test")  # foo() takes 2 positional arguments but 3 were given
+r6 = foo(1, 2.0, b=3.0)  # foo() got multiple values for argument 'b'
+r7 = foo(
+    a=1
+)  # foo() got some positional-only arguments passed as keyword arguments: 'a'
+r8 = foo(g="test")  # foo() got an unexpected keyword argument 'g'
+
+r9a = foo(1, 2.0, c="test")
+r9b = foo(1, b=2.0, c="test")
+r9c = foo(1, c="test", b=2.0)
+
+r10 = foo("a", 3, c=False)  # wrong argument types

examples/01_simple_type_checking/06_overloads.midas

@@ -0,0 +1,10 @@
+type T1 = object
+type T2 = object
+type Foo = object
+type T2b = T2
+
+extend Foo {
+    def bar: fn(T1, /) -> int
+    def bar: fn(T2, /) -> float
+    def bar: fn(T2b, /) -> int
+}

examples/01_simple_type_checking/06_overloads.py

@@ -0,0 +1,18 @@
+# type: ignore
+# ruff: disable [F821]
+
+foo: Foo
+t1: T1
+t2: T2
+
+a = foo.bar(t1)
+b = foo.bar(t2)
+
+func = foo.bar
+
+c = func(t1)
+d = func(t2)
+
+t2b: T2b
+
+e = foo.bar(t2b)

midas/checker/python.py

@@ -12,12 +12,16 @@ from midas.checker.reporter import FileReporter, Reporter
 from midas.checker.resolver import Resolver
 from midas.checker.types import (
     Function,
+    OverloadedFunction,
     Type,
     UnitType,
     UnknownType,
+    unfold_type,
 )
 from midas.parser.python import PythonParser
 
+TypedExpr = tuple[p.Expr, Type]
+
 
 class ReturnException(Exception):
     pass
@@ -30,6 +34,12 @@ class MappedArgument:
     argument: Function.Argument
 
 
+@dataclass(frozen=True, kw_only=True)
+class OverloadCandidate:
+    function: Function
+    mapped: list[MappedArgument]
+
+
 class PythonTyper(
     p.Stmt.Visitor[None],
     p.Expr.Visitor[Type],
@@ -354,26 +364,7 @@ class PythonTyper(
             )
             return UnknownType()
 
-        match operation:
-            case Function() as function:
-                if not self._check_arity(function, 1, 0, 0):
-                    self.reporter.error(
-                        location,
-                        f"Wrong definition of binary operation. Expected function with 1 positional-only parameters, got {function}",
-                    )
-                    return UnknownType()
-
-                rhs: Function.Argument = function.pos_args[0]
-                if not self.is_subtype(right, rhs.type):
-                    self.reporter.error(
-                        location,
-                        f"Wrong type for right-hand side, expected {rhs.type}, got {right}",
-                    )
-                    return UnknownType()
-                return function.returns
-            case _:
-                self.reporter.warning(location, f"Unsupported operation {operation}")
-                return UnknownType()
+        return self._get_call_result(location, operation, [(right_expr, right)], {})
 
     def visit_unary_expr(self, expr: p.UnaryExpr) -> Type:
         method: Optional[str] = UNARY_METHODS.get(expr.operator.__class__)
@@ -393,35 +384,24 @@ class PythonTyper(
             )
             return UnknownType()
 
-        match operation:
-            case Function() as function:
-                if not self._check_arity(function, 0, 0, 0):
-                    self.reporter.error(
-                        expr.location,
-                        f"Wrong definition of unary operation. Expected function with 0 parameters, got {function}",
+        return self._get_call_result(
+            expr.location, operation, [(expr.right, operand)], {}
         )
-                    return UnknownType()
-                return function.returns
-            case _:
-                self.reporter.warning(
-                    expr.location, f"Unsupported operation {operation}"
-                )
-                return UnknownType()
 
     def visit_call_expr(self, expr: p.CallExpr) -> Type:
         callee: Type = self.type_of(expr.callee)
-        if not isinstance(callee, Function):
-            self.reporter.error(expr.callee.location, "Callee is not a function")
-            return UnknownType()
-        function: Function = callee
-        mapped: list[MappedArgument] = self.map_call_arguments(function, expr)
-        for arg in mapped:
-            if not self.is_subtype(arg.type, arg.argument.type):
-                self.reporter.error(
-                    arg.expr.location,
-                    f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
+        positional: list[TypedExpr] = [
+            (arg, self.type_of(arg)) for arg in expr.arguments
+        ]
+        keywords: dict[str, TypedExpr] = {
+            name: (arg, self.type_of(arg)) for name, arg in expr.keywords.items()
+        }
+        return self._get_call_result(
+            location=expr.location,
+            callee=callee,
+            positional=positional,
+            keywords=keywords,
         )
-        return function.returns
 
     def visit_get_expr(self, expr: p.GetExpr) -> Type:
         object: Type = self.type_of(expr.object)
@@ -523,29 +503,9 @@ class PythonTyper(
             return UnknownType()
 
         index: Type = self.type_of(expr.index)
-
-        match operation:
-            case Function() as function:
-                if not self._check_arity(function, 1, 0, 0):
-                    self.reporter.error(
-                        expr.location,
-                        f"Wrong definition of __getitem__. Expected function with 1 positional-only parameters, got {function}",
+        return self._get_call_result(
+            expr.location, operation, [(expr.index, index)], {}
         )
-                    return UnknownType()
-
-                index_arg: Function.Argument = function.pos_args[0]
-                if not self.is_subtype(index, index_arg.type):
-                    self.reporter.error(
-                        expr.location,
-                        f"Wrong index type, expected {index_arg.type}, got {index}",
-                    )
-                    return UnknownType()
-                return function.returns
-            case _:
-                self.reporter.warning(
-                    expr.location, f"Unsupported operation {operation}"
-                )
-                return UnknownType()
 
     def visit_base_type(self, node: p.BaseType) -> Type:
         base: Type
@@ -572,29 +532,192 @@ class PythonTyper(
         self.reporter.warning(node.location, "FrameType not yet supported")
         return UnknownType()
 
+    def _get_call_result(
+        self,
+        location: Location,
+        callee: Type,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Type:
+        """Get the result type of a function call
+
+        If the function has overloads, the function will try to resolve the
+        appropriate signature.
+        Argument types are matched to the defined parameters.
+        The function doesn't take the raw expression as a parameter to accomodate
+        for desugared calls such as for operators.
+
+        Args:
+            location (Location): the call location
+            callee (Type): the called function
+            positional (list[TypedExpr]): the list positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+
+        Returns:
+            Type: the return type of the call, or `UnknownType` if either
+            the call is invalid or no overload matched the arguments uniquely
+        """
+        match callee:
+            case Function() as function:
+                valid: bool
+                mapped: list[MappedArgument]
+                valid, mapped = self.map_call_arguments(
+                    function, location, positional, keywords
+                )
+                valid = valid and self._are_arguments_valid(mapped)
+                if not valid:
+                    return UnknownType()
+                return function.returns
+
+            case OverloadedFunction(overloads=overloads):
+                function = self._match_overload(
+                    overloads, location, positional, keywords
+                )
+                if function is None:
+                    return UnknownType()
+                return function.returns
+            case _:
+                self.reporter.error(location, f"{callee} is not callable")
+                return UnknownType()
+
+    def _are_arguments_valid(
+        self,
+        arguments: list[MappedArgument],
+        report_errors: bool = True,
+    ) -> bool:
+        """Check whether the passed argument types correspond to their matched parameter definitions
+
+        Args:
+            arguments (list[MappedArgument]): the list of argument/parameter pairs
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
+
+        Returns:
+            bool: True if all arguments fit the matching parameter definitions, False otherwise
+        """
+        valid: bool = True
+        for arg in arguments:
+            if not self.is_subtype(arg.type, arg.argument.type):
+                if report_errors:
+                    self.reporter.error(
+                        arg.expr.location,
+                        f"Wrong type for argument '{arg.argument.name}', expected {arg.argument.type}, got {arg.type}",
+                    )
+                valid = False
+        return valid
+
+    def _match_overload(
+        self,
+        overloads: list[Type],
+        location: Location,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+    ) -> Optional[Function]:
+        """Try and resolve the appropriate overload for the given arguments
+
+        Args:
+            overloads (list[Type]): the list of possible overloads
+            location (Location): the call location
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keywords arguments
+
+        Returns:
+            Optional[Function]: the resolved function signature if it can be
+            determined unambigously, or `None`.
+        """
+        candidates: list[OverloadCandidate] = []
+        for overload in overloads:
+            function: Type = unfold_type(overload)
+            if not isinstance(function, Function):
+                self.logger.error(
+                    f"Overload is not a function: {overload} is {function}"
+                )
+                continue
+            valid, mapped = self.map_call_arguments(
+                function=function,
+                location=location,
+                positional=positional,
+                keywords=keywords,
+                report_errors=False,
+            )
+            if valid and self._are_arguments_valid(mapped, report_errors=False):
+                candidates.append(
+                    OverloadCandidate(
+                        function=function,
+                        mapped=mapped,
+                    )
+                )
+
+        pos_types: str = ", ".join(str(type) for _, type in positional)
+        kw_types: str = ", ".join(
+            f"{name}: {type}" for name, (_, type) in keywords.items()
+        )
+        for_args: str = f"for arguments pos=[{pos_types}] and kw={{{kw_types}}}"
+
+        n_candidates: int = len(candidates)
+
+        # Exactly 1 match -> return it
+        if n_candidates == 1:
+            return candidates[0].function
+
+        # No match -> invalid call
+        if n_candidates == 0:
+            self.reporter.error(
+                location,
+                f"No matching overload in {overloads} {for_args}",
+            )
+            return None
+
+        # Multiple matches -> see if one <: all others (more specific)
+        for i1, c1 in enumerate(candidates):
+            mapped1: list[MappedArgument] = c1.mapped
+            best_match: bool = True
+            for i2, c2 in enumerate(candidates):
+                if i1 == i2:
+                    continue
+                mapped2: list[MappedArgument] = c2.mapped
+                if not self._are_mapped_subtypes(mapped1, mapped2):
+                    best_match = False
+                    break
+                self.logger.debug(f"{c1.function} is a full overload of {c2.function}")
+            if best_match:
+                return c1.function
+
+        candidates_str: str = ", ".join(
+            str(candidate.function) for candidate in candidates
+        )
+        self.reporter.error(
+            location,
+            f"Multiple matching overloads {for_args}: {candidates_str}",
+        )
+        return None
+
     def map_call_arguments(
-        self, function: Function, call: p.CallExpr
-    ) -> list[MappedArgument]:
-        """Map call arguments to function parameters as defined in its signature
+        self,
+        function: Function,
+        location: Location,
+        positional: list[TypedExpr],
+        keywords: dict[str, TypedExpr],
+        report_errors: bool = True,
+    ) -> tuple[bool, list[MappedArgument]]:
+        """Map call arguments to a function's parameters as defined in its signature
 
         This method maps positional-only, keyword-only and mixed parameter definitions
         with the arguments passed at the call site
 
-        Any mismatched, missing or unexpected argument is reported as a diagnostic
+        Any mismatched, missing or unexpected argument is reported as a diagnostic,
+        unless `report_errors` is set to `False`
 
         Args:
             function (Function): the function definition
-            call (p.CallExpr): the call expression
+            location (Location): the call location
+            positional (list[TypedExpr]): the list of positional arguments
+            keywords (dict[str, TypedExpr]): the map of keyword arguments
+            report_errors (bool, optional): whether type errors should be reported as diagnostics. Defaults to True.
 
         Returns:
-            list[MappedArgument]: the list of mapped arguments
+            tuple[bool, list[MappedArgument]]: a boolean reporting whether
+            the call is valid and the list of mapped arguments
         """
-        positional: list[tuple[p.Expr, Type]] = [
-            (arg, self.type_of(arg)) for arg in call.arguments
-        ]
-        keywords: dict[str, tuple[p.Expr, Type]] = {
-            name: (arg, self.type_of(arg)) for name, arg in call.keywords.items()
-        }
         set_args: set[str] = set()
 
         required_positional: list[str] = [
@@ -612,6 +735,8 @@ class PythonTyper(
             arg.name: arg for arg in function.kw_args
         }
 
+        valid_call: bool = True
+
         # TODO: handle *args and **kwargs sinks
         for arg in positional:
             param: Function.Argument
@@ -620,7 +745,11 @@ class PythonTyper(
             elif len(mixed_params) != 0:
                 param = mixed_params.pop(0)
             else:
-                self.reporter.error(arg[0].location, "Too many positional arguments")
+                if report_errors:
+                    self.reporter.error(
+                        arg[0].location, "Too many positional arguments"
+                    )
+                valid_call = False
                 break
             name: str = param.name
             if name in required_positional:
@@ -640,6 +769,7 @@ class PythonTyper(
         for name, arg in keywords.items():
             param: Function.Argument
             if name not in kw_params:
+                if report_errors:
                     if name in set_args:
                         self.reporter.error(
                             arg[0].location, f"Multiple values for argument '{name}'"
@@ -648,6 +778,7 @@ class PythonTyper(
                         self.reporter.error(
                             arg[0].location, f"Unknown keyword argument '{name}'"
                         )
+                valid_call = False
                 continue
             param = kw_params.pop(name)
             if name in required_positional:
@@ -674,32 +805,51 @@ class PythonTyper(
         if len(required_positional) != 0:
             plural: str = "" if len(required_positional) == 1 else "s"
             args: str = join_args(required_positional)
+            if report_errors:
                 self.reporter.error(
-                call.location,
+                    location,
                     f"Missing required positional argument{plural}: {args}",
                 )
+            valid_call = False
 
         if len(required_keyword) != 0:
             plural: str = "" if len(required_keyword) == 1 else "s"
             args: str = join_args(required_keyword)
+            if report_errors:
                 self.reporter.error(
-                call.location,
+                    location,
                     f"Missing required keyword argument{plural}: {args}",
                 )
+            valid_call = False
 
-        return mapped
+        return valid_call, mapped
 
-    def _check_arity(
-        self,
-        function: Function,
-        n_pos: Optional[int] = None,
-        n_mixed: Optional[int] = None,
-        n_keyword: Optional[int] = None,
+    def _are_mapped_subtypes(
+        self, mapped1: list[MappedArgument], mapped2: list[MappedArgument]
     ) -> bool:
-        if n_pos is not None and len(function.pos_args) != n_pos:
-            return False
-        if n_mixed is not None and len(function.args) != n_mixed:
-            return False
-        if n_keyword is not None and len(function.kw_args) != n_keyword:
+        """Check whether the given argument mappings are subtype/supertype of one another
+
+        This function checks whether the argument mappings `mapped1` are subtypes
+        of `mapped2`. If any of the parameter type in `mapped1` is not a subtype
+        of the corresponding parameter in `mapped2`, `False` is returned.
+
+        This is used to check whether a given overload is
+        a more specific function/ a subtype of another.
+
+        Args:
+            mapped1 (list[MappedArgument]): the first argument mappings (subtype)
+            mapped2 (list[MappedArgument]): the second argument mappings (supertype)
+
+        Returns:
+            bool: `True` if `mapped1` is a subtype of `mapped2`, `False` otherwise
+        """
+        by_expr: dict[p.Expr, Type] = {}
+        for arg in mapped1:
+            by_expr[arg.expr] = arg.argument.type
+
+        for arg in mapped2:
+            type2: Type = arg.argument.type
+            type1: Type = by_expr[arg.expr]
+            if not self.is_subtype(type1, type2):
                 return False
         return True