Support overloading with TypedDict (#3612)

Fix #3609.

Author: JukkaL

mypy/checkexpr.py

@@ -2675,6 +2675,12 @@ def overload_arg_similarity(actual: Type, formal: Type) -> int:
             return overload_arg_similarity(actual.ret_type, formal.item)
         else:
             return 0
+    if isinstance(actual, TypedDictType):
+        if isinstance(formal, TypedDictType):
+            # Don't support overloading based on the keys or value types of a TypedDict since
+            # that would be complicated and probably only marginally useful.
+            return 2
+        return overload_arg_similarity(actual.fallback, formal)
     if isinstance(formal, Instance):
         if isinstance(actual, CallableType):
             actual = actual.fallback

mypy/meet.py

@@ -84,6 +84,10 @@ class C(A, B): ...
         t = t.erase_to_union_or_bound()
     if isinstance(s, TypeVarType):
         s = s.erase_to_union_or_bound()
+    if isinstance(t, TypedDictType):
+        t = t.as_anonymous().fallback
+    if isinstance(s, TypedDictType):
+        s = s.as_anonymous().fallback
     if isinstance(t, Instance):
         if isinstance(s, Instance):
             # Consider two classes non-disjoint if one is included in the mro

test-data/unit/check-typeddict.test

@@ -1086,6 +1086,125 @@ X = TypedDict('Y', {'x': int})  # E: First argument 'Y' to TypedDict() does not
 [builtins fixtures/dict.pyi]
 
 
+-- Overloading
+
+[case testTypedDictOverloading]
+from typing import overload, Iterable
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+
+@overload
+def f(x: Iterable[str]) -> str: ...
+@overload
+def f(x: int) -> int: ...
+def f(x): pass
+
+a: A
+reveal_type(f(a))  # E: Revealed type is 'builtins.str'
+reveal_type(f(1))  # E: Revealed type is 'builtins.int'
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testTypedDictOverloading2]
+from typing import overload, Iterable
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+
+@overload
+def f(x: Iterable[int]) -> None: ...
+@overload
+def f(x: int) -> None: ...
+def f(x): pass
+
+a: A
+f(a)  # E: Argument 1 to "f" has incompatible type "A"; expected Iterable[int]
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testTypedDictOverloading3]
+from typing import overload
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+
+@overload
+def f(x: str) -> None: ...
+@overload
+def f(x: int) -> None: ...
+def f(x): pass
+
+a: A
+f(a)  # E: No overload variant of "f" matches argument types [TypedDict(x=builtins.int, _fallback=__main__.A)]
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testTypedDictOverloading4]
+from typing import overload
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+B = TypedDict('B', {'x': str})
+
+@overload
+def f(x: A) -> int: ...
+@overload
+def f(x: int) -> str: ...
+def f(x): pass
+
+a: A
+b: B
+reveal_type(f(a)) # E: Revealed type is 'builtins.int'
+reveal_type(f(1)) # E: Revealed type is 'builtins.str'
+f(b) # E: Argument 1 to "f" has incompatible type "B"; expected "A"
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testTypedDictOverloading5]
+from typing import overload
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+B = TypedDict('B', {'y': str})
+C = TypedDict('C', {'y': int})
+
+@overload
+def f(x: A) -> None: ...
+@overload
+def f(x: B) -> None: ...
+def f(x): pass
+
+a: A
+b: B
+c: C
+f(a)
+f(b)
+f(c) # E: Argument 1 to "f" has incompatible type "C"; expected "A"
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testTypedDictOverloading6]
+from typing import overload
+from mypy_extensions import TypedDict
+
+A = TypedDict('A', {'x': int})
+B = TypedDict('B', {'y': str})
+
+@overload
+def f(x: A) -> int: ... # E: Overloaded function signatures 1 and 2 overlap with incompatible return types
+@overload
+def f(x: B) -> str: ...
+def f(x): pass
+
+a: A
+b: B
+reveal_type(f(a)) # E: Revealed type is 'Any'
+reveal_type(f(b)) # E: Revealed type is 'Any'
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+
 -- Special cases
 
 [case testForwardReferenceInTypedDict]