Merge pull request #186 from pradeep90/master

Shrink heterogenous list types to avoid rewriting TypedDict to Dict.
Instagram · Jun 4, 2020 · 330dc6f · 330dc6f
2 parents 5e65f1b + b0b6c11
commit 330dc6f
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 3 deletions.
diff --git a/monkeytype/typing.py b/monkeytype/typing.py
@@ -50,6 +50,10 @@
 # file live in typing.pyi.
 
 
+def is_list(typ: type) -> bool:
+    return is_generic(typ) and name_of_generic(typ) == 'List'
+
+
 def make_typed_dict(*, required_fields=None, optional_fields=None) -> type:
     required_fields = required_fields or {}
     optional_fields = optional_fields or {}
@@ -119,9 +123,16 @@ def shrink_types(types, max_typed_dict_size):
     if all(is_anonymous_typed_dict(typ) for typ in types):
         return shrink_typed_dict_types(types, max_typed_dict_size)
     # Don't rewrite anonymous TypedDict to Dict if the types are all the same,
-    # such as List[TypedDict(...)].
+    # such as [Tuple[TypedDict(...)], Tuple[TypedDict(...)]].
     if all(types_equal(typ, types[0]) for typ in types[1:]):
         return types[0]
+
+    # If they are all lists, shrink their argument types. This way, we avoid
+    # rewriting heterogenous anonymous TypedDicts to Dict.
+    if all(is_list(typ) for typ in types):
+        annotation = shrink_types((getattr(typ, '__args__')[0] for typ in types), max_typed_dict_size)
+        return List[annotation]
+
     all_dict_types = tuple(RewriteAnonymousTypedDictToDict().rewrite(typ) for typ in types)
     return Union[all_dict_types]
 

diff --git a/tests/test_typing.py b/tests/test_typing.py
@@ -30,6 +30,7 @@
     RewriteAnonymousTypedDictToDict,
     field_annotations,
     get_type,
+    is_list,
     is_typed_dict,
     make_typed_dict,
     shrink_types,
@@ -101,6 +102,17 @@ class TestTypesEqual:
     def test_types_equal(self, typ, other_type, expected_output):
         assert (types_equal(typ, other_type) == expected_output)
 
+    @pytest.mark.parametrize(
+        'typ, expected',
+        [
+            (List[int], True),
+            (typing_Tuple[int], False),
+            (int, False),
+        ],
+    )
+    def test_is_list(self, typ, expected):
+        assert is_list(typ) == expected
+
 
 class TestMakeTypedDict:
     @pytest.mark.parametrize(
@@ -315,13 +327,41 @@ def test_shrink_non_uniform_typed_dict_types(self, types, expected_type):
                 ),
                 List[Dict[str, int]],
             ),
-            # Same. We don't currently shrink the inner types even if the outer types are the same.
             (
                 (
                     List[make_typed_dict(required_fields={'a': int})],
                     List[make_typed_dict(required_fields={'b': int})],
                 ),
-                List[Dict[str, int]],
+                List[make_typed_dict(optional_fields={'a': int, 'b': int})],
+            ),
+            (
+                (
+                    make_typed_dict(required_fields={"foo": List[make_typed_dict(required_fields={'a': int})]}),
+                    make_typed_dict(required_fields={"foo": List[make_typed_dict(required_fields={'a': int})]}),
+                ),
+                make_typed_dict(required_fields={"foo": List[make_typed_dict(required_fields={'a': int})]}),
+            ),
+            (
+                (
+                    make_typed_dict(required_fields={"foo": List[make_typed_dict(required_fields={'a': int})]}),
+                    make_typed_dict(required_fields={"foo": List[make_typed_dict(required_fields={'b': int})]}),
+                ),
+                make_typed_dict(required_fields={"foo": List[make_typed_dict(optional_fields={'a': int, 'b': int})]}),
+            ),
+            (
+                (
+                    typing_Tuple[make_typed_dict(required_fields={'a': int})],
+                    typing_Tuple[make_typed_dict(required_fields={'a': int})],
+                ),
+                typing_Tuple[make_typed_dict(required_fields={'a': int})],
+            ),
+            # We don't currently shrink the inner types for Tuples.
+            (
+                (
+                    typing_Tuple[make_typed_dict(required_fields={'a': int})],
+                    typing_Tuple[make_typed_dict(required_fields={'b': int})],
+                ),
+                typing_Tuple[Dict[str, int]],
             ),
             # Fall back to Dict when the resulting TypedDict would be too large.
             # Keep any nested anonymous TypedDicts, though.