Fix type hints for MutableMap.get()

Summary: Fix type hints for `MutableMap.get()` and add typing test for `get()` and `__setitem__()`.

Reviewed By: Filip-F

Differential Revision: D66253483

fbshipit-source-id: dbdd7a2522d8b5b82ec611da67f88c6dec5b9955

third-party/thrift/src/thrift/lib/python/mutable_containers.pyi

@@ -213,8 +213,10 @@ class MutableMap(MutableMappingAbc[K, V]):
     def __len__(self) -> int: ...
     def __getitem__(self, key: K) -> V: ...
     def __iter__(self) -> ValueIterator[K]: ...
-    # pyre-ignore[14]: Inconsistent override
-    def get(self, key: K, default: Optional[V] = None) -> V: ...
+    @overload
+    def get(self, key: K, /) -> V | None: ...
+    @overload
+    def get(self, key: K, /, default: V | _T) -> V | _T: ...
     @overload
     def __setitem__(self, key: K, value: V) -> None: ...
     @overload

third-party/thrift/src/thrift/lib/python/test/mutable_map_test.py

@@ -19,6 +19,8 @@
 import string
 import unittest
 
+from typing import cast, Optional
+
 from thrift.python.mutable_containers import (
     MapItemsView,
     MapKeysView,
@@ -42,6 +44,107 @@ def _create_MutableMap_str_i32(map: dict[str, int]) -> MutableMap[str, int]:
     return MutableMap(typeinfo_string, typeinfo_i32, map)
 
 
+class MutableMapTypeHints(unittest.TestCase):
+    """
+    Tests type hints for MutableMap class.
+
+    This test suite checks that Pyre correctly identifies and raises errors
+    when there are type mismatches in function calls or assignments.
+
+    Casting an object to MutableMap[K, V] is sufficient to check the type hints.
+
+    The `pyre-ignore` comments in the test are intended to suppress real errors,
+    as we run Pyre with the `--report-unused-ignores` flag.
+    """
+
+    def test_type_hints(self) -> None:
+        mutable_map = cast(MutableMap[str, int], object())
+        try:
+            ### get() ###
+            v1: Optional[int] = mutable_map.get("key")
+
+            # pyre-ignore[9]: v2 is type `int` but is used as type `Optional[int]`
+            v2: int = mutable_map.get("key")
+
+            v3: int = mutable_map.get("key", 42)
+            v4: int | str = mutable_map.get("key", "value")
+
+            # pyre-ignore[6]: 1st positional argument, expected `str` but got `int`
+            v5: int = mutable_map.get(999, 42)
+
+            # to silence F841: not used variable
+            _ = (v1, v2, v3, v4, v5)
+
+            ###################################################################
+
+            ### __setitem__() ###
+            mutable_map["key"] = 1
+
+            # pyre-ignore[6]: expected `int` but got `str`
+            mutable_map["key"] = "value"
+            # pyre-ignore[6]: expected `int` but got `float`
+            mutable_map["key"] = 1.0
+
+            # pyre-ignore[6]: expected `str` but got `int`
+            mutable_map[1] = 1
+            # pyre-ignore[6]: expected `str` but got `bytes`
+            mutable_map[b"key"] = 1
+
+            # pyre-ignore[6]: expected `str` but got `int` and expected `int` but got `str`
+            mutable_map[1] = "value"
+
+        except Exception:
+            pass
+
+    def test_type_hints_with_container_value(self) -> None:
+        mutable_map = cast(MutableMap[str, MutableList[int]], object())
+        try:
+            ### get() ###
+            v1: Optional[MutableList[int]] = mutable_map.get("key")
+
+            # pyre-ignore[9]: v2 is type `MutableList[int]` but is used as type `Optional[MutableList[int]]`.
+            v2: MutableList[int] = mutable_map.get("key")
+
+            default_value_1 = cast(MutableList[int], object())
+            v3: MutableList[int] = mutable_map.get("key", default_value_1)
+
+            default_value_2 = cast(MutableList[str], object())
+            v4: MutableList[int] | MutableList[str] = mutable_map.get(
+                "key", default_value_2
+            )
+
+            # pyre-ignore[6]: 1st positional argument, expected `str` but got `int`
+            v5: MutableList[int] = mutable_map.get(999, default_value_1)
+
+            # to silence F841: not used variable
+            _ = (v1, v2, v3, v4, v5)
+
+            ###################################################################
+
+            ### __setitem__() ###
+            mutable_list_int = cast(MutableList[int], object())
+            mutable_list_str = cast(MutableList[str], object())
+
+            mutable_map["key"] = to_thrift_list([])
+            mutable_map["key"] = mutable_list_int
+
+            # pyre-ignore[6]: expected `MutableList[int]` but got `MutableList[str]`
+            mutable_map["key"] = mutable_list_str
+            # pyre-ignore[6]: expected `MutableList[int]` but got `List[int]`
+            mutable_map["key"] = [1, 2, 3]
+
+            # pyre-ignore[6]: expected `str` but got `int`
+            mutable_map[1] = mutable_list_int
+            # pyre-ignore[6]: expected `str` but got `bytes`
+            mutable_map[b"key"] = mutable_list_int
+
+            # pyre-ignore[6]: expected `str` but got `int` and expected `MutableList[int]` but got `MutableList[str]`
+            mutable_map[1] = mutable_list_str
+
+        except Exception:
+            pass
+
+
 class MutableMapTest(unittest.TestCase):
     def test_smoke(self) -> None:
         mutable_map = MutableMap(typeinfo_string, typeinfo_i32, {})

third-party/thrift/src/thrift/test/thrift-python/struct_test.py

@@ -1218,7 +1218,6 @@ def test_create_and_assign_for_map(self) -> None:
         self.assertEqual(None, s.unqualified_map_string_i32.get("Not Exists"))
         self.assertEqual(
             "MyDefaultValue",
-            # pyre-ignore[6]: Intentional for test
             s.unqualified_map_string_i32.get("Not Exists", "MyDefaultValue"),
         )