Implement Overloaded. (#92)

Author: dnwpark

tests/test_type_eval.py

@@ -16,6 +16,7 @@
     TypeVar,
     Union,
     get_args,
+    overload,
 )
 
 import pytest
@@ -45,6 +46,7 @@
     Member,
     Members,
     NewProtocol,
+    Overloaded,
     Param,
     Slice,
     SpecialFormEllipsis,
@@ -466,6 +468,46 @@ def f[T](self, x: T) -> OnlyIntToSet[T]: ...
     )
 
 
+def test_getmember_04():
+    class C:
+        @overload
+        def f(self, x: int) -> set[int]: ...
+
+        @overload
+        def f[T](self, x: T) -> T: ...
+
+        def f(self, x): ...
+
+    m = eval_typing(GetMember[C, Literal["f"]])
+    mt = eval_typing(GetType[m])
+    assert mt.__origin__ is Overloaded
+    assert len(mt.__args__) == 2
+
+    # Non-generic overload
+    assert (
+        eval_typing(IsAssignable[GetArg[mt, Overloaded, Literal[0]], Callable])
+        == _BoolLiteral[True]
+    )
+    assert (
+        mt.__args__[0]
+        == Callable[
+            [Param[Literal["self"], C], Param[Literal["x"], int]], set[int]
+        ]
+    )
+
+    # Generic overload
+    assert (
+        eval_typing(
+            IsAssignable[GetArg[mt, Overloaded, Literal[1]], GenericCallable]
+        )
+        == _BoolLiteral[True]
+    )
+    assert (
+        eval_typing(mt.__args__[1].__args__[1](int))
+        == Callable[[Param[Literal["self"], C], Param[Literal["x"], int]], int]
+    )
+
+
 def test_getarg_never():
     d = eval_typing(GetArg[Never, object, Literal[0]])
     assert d is Never

typemap/type_eval/_apply_generic.py

@@ -291,7 +291,14 @@ def _resolved_function_signature(func, args):
     return sig
 
 
-def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
+def get_local_defns(
+    boxed: Boxed,
+) -> tuple[
+    dict[str, Any],
+    dict[
+        str, types.FunctionType | classmethod | staticmethod | WrappedOverloads
+    ],
+]:
     from typemap.typing import GenericCallable
 
     annos: dict[str, Any] = {}
@@ -327,6 +334,8 @@ def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
                 # XXX: This is totally wrong; we still need to do
                 # substitute in class vars
                 local_fn = stuff
+            elif overloads := typing.get_overloads(stuff):
+                local_fn = WrappedOverloads(tuple(overloads))
 
             # If we got stuck, we build a GenericCallable that
             # computes the type once it has been given type
@@ -370,6 +379,11 @@ def lam(*vs):
     return annos, dct
 
 
+@dataclasses.dataclass(frozen=True)
+class WrappedOverloads:
+    functions: tuple[typing.Callable[..., Any], ...]
+
+
 def flatten_class_new_proto(cls: type) -> type:
     # This is a hacky version of flatten_class that works by using
     # NewProtocol on Members!

typemap/type_eval/_eval_operators.py

@@ -39,6 +39,7 @@
     Member,
     Members,
     NewProtocol,
+    Overloaded,
     Param,
     RaiseError,
     Slice,
@@ -169,6 +170,17 @@ def get_annotated_method_hints(cls, *, ctx):
                     object,
                     acls,
                 )
+            elif isinstance(attr, _apply_generic.WrappedOverloads):
+                overloads = [
+                    _function_type(_eval_types(of, ctx), receiver_type=acls)
+                    for of in attr.functions
+                ]
+                hints[name] = (
+                    Overloaded[*overloads],
+                    ("ClassVar",),
+                    object,
+                    acls,
+                )
 
     return hints
 

typemap/typing.py

@@ -101,6 +101,10 @@ def __class_getitem__(cls, params):
         return _GenericCallableGenericAlias(cls, (typevars, func))
 
 
+class Overloaded[*Callables]:
+    pass
+
+
 ###