feat(overloads): Expand finite sum types into a union of possible types

Author: mimre25

mypy/checkexpr.py

@@ -2709,6 +2709,16 @@ def check_overload_call(
         # Normalize unpacked kwargs before checking the call.
         callee = callee.with_unpacked_kwargs()
         arg_types = self.infer_arg_types_in_empty_context(args)
+
+        # Expand finite sum types into unions
+        # See https://github.com/python/mypy/issues/14764#issuecomment-3054510950
+        # And https://typing.python.org/en/latest/spec/overload.html#argument-type-expansion
+        arg_types = [
+                try_expanding_sum_type_to_union(arg_type,  arg_type.type.fullname)
+                if isinstance(arg_type, Instance) else arg_type
+                for arg_type in arg_types
+                ]
+
         # Step 1: Filter call targets to remove ones where the argument counts don't match
         plausible_targets = self.plausible_overload_call_targets(
             arg_types, arg_kinds, arg_names, callee

mypy/typeops.py

@@ -1050,7 +1050,12 @@ class Status(Enum):
         ]
         return make_simplified_union(items, contract_literals=False)
 
+
     if isinstance(typ, Instance) and typ.type.fullname == target_fullname:
+        if isinstance(typ.last_known_value, LiteralType):
+            # fallback for Literal[True] and Literal[False]
+            return typ
+
         if typ.type.fullname == "builtins.bool":
             items = [LiteralType(True, typ), LiteralType(False, typ)]
             return make_simplified_union(items, contract_literals=False)

test-data/unit/check-overloading.test

@@ -5266,6 +5266,47 @@ tmp/lib.pyi:3: error: Name "func" already defined on line 1
 tmp/lib.pyi:3: error: Name "overload" is not defined
 main:3: note: Revealed type is "Any"
 
+[case testOverloadCheckExpandsBools]
+from typing import Literal, overload, Union
+
+@overload
+def foo(x: Literal[False]) -> None: ...
+@overload
+def foo(x: Literal[True]) -> int: ...
+
+def foo(x: bool) -> Union[None, int]: ...
+
+reveal_type(foo(True)) # N: Revealed type is "builtins.int"
+reveal_type(foo(False)) # N: Revealed type is "None"
+x: bool
+reveal_type(foo(x)) # N: Revealed type is "Union[builtins.int, None]"
+
+[case testOverloadCheckExpandsEnums]
+from typing import Literal, overload, Union
+import enum
+
+class Color(enum.Enum):
+    RED = 1
+    BLUE = 2
+    YELLOW = 3
+
+@overload
+def foo(x: Literal[Color.RED]) -> None: ...
+@overload
+def foo(x: Literal[Color.BLUE]) -> int: ...
+@overload
+def foo(x: Literal[Color.YELLOW]) -> str: ...
+
+def foo(x: Color) -> Union[None, int, str]: ...
+
+reveal_type(foo(Color.RED)) # N: Revealed type is "None"
+reveal_type(foo(Color.BLUE)) # N: Revealed type is "builtins.int"
+reveal_type(foo(Color.YELLOW)) # N: Revealed type is "builtins.str"
+
+x: Color
+reveal_type(foo(x)) # N: Revealed type is "Union[None, builtins.int, builtins.str]"
+[builtins fixtures/tuple.pyi]
+
 [case testLiteralSubtypeOverlap]
 from typing import Literal, overload