Trace and shrink dictionaries as TypedDicts if possible.

* Trace a dictionary as an anonymous TypedDict if all keys are strings (up to a limit).
* Conservatively shrink multiple TypedDicts as one only if they are all identical. Otherwise, fall back to the default behaviour.
* Generate a TypedDict class definition stub at the top of the module stub and name it using the parameter name.
* Add MAX_TYPED_DICT_SIZE field to Config.
* Hand-roll an equality test because the default equality test for TypedDict fails for identical invocations: TypedDict('Foo', {'a': int}) != TypedDict('Foo', {'a': int}).
* Trace and shrink return types and yield types as well.
* Handle nested dictionaries.

* Add `make_forward_ref` so that we can use ForwardRefs. I deleted a test for a nested `Dict[str, TypedDict]` because it kept failing on Python 3.6.

Pipfile

@@ -1,22 +1,22 @@
 [[source]]
-
 url = "https://pypi.python.org/simple"
 verify_ssl = true
 name = "pypi"
 
 [dev-packages]
-
 "flake8" = "*"
 pytest = "*"
 mypy = "*"
 pytest-smartcov = "*"
-"e1839a8" = {path = ".", editable = true}
+e1839a8 = {path = ".",editable = true}
 cython = "*"
 sphinx = "*"
 twine = "*"
 django = "*"
 tox = "*"
 
 [pipenv]
-
 keep_outdated = true
+
+[packages]
+stringcase = "*"

monkeytype/__init__.py

@@ -29,4 +29,5 @@ def trace(config: Optional[Config] = None) -> ContextManager:
         logger=config.trace_logger(),
         code_filter=config.code_filter(),
         sample_rate=config.sample_rate(),
+        max_typed_dict_size=config.max_typed_dict_size(),
     )

monkeytype/compat.py

@@ -30,6 +30,13 @@ def name_of_generic(typ: Any) -> str:
     def is_forward_ref(typ: Any) -> bool:
         return isinstance(typ, ForwardRef)
 
+    def make_forward_ref(s: str) -> type:
+        return ForwardRef(s)
+
+    def repr_forward_ref() -> str:
+        """For checking the test output when ForwardRef is printed."""
+        return 'ForwardRef'
+
 except ImportError:
     # Python 3.6
     from typing import _Any, _Union, GenericMeta, _ForwardRef  # type: ignore
@@ -54,3 +61,10 @@ def name_of_generic(typ: Any) -> str:
 
     def is_forward_ref(typ: Any) -> bool:
         return isinstance(typ, _ForwardRef)
+
+    def make_forward_ref(s: str) -> type:
+        return _ForwardRef(s)
+
+    def repr_forward_ref() -> str:
+        """For checking the test output when ForwardRef is printed."""
+        return '_ForwardRef'

monkeytype/config.py

@@ -85,6 +85,11 @@ def query_limit(self) -> int:
         """Maximum number of traces to query from the call trace store."""
         return 2000
 
+    @abstractmethod
+    def max_typed_dict_size(self) -> int:
+        """Size up to which a dictionary will be traced as a TypedDict."""
+        pass
+
 
 lib_paths = {sysconfig.get_path(n) for n in ['stdlib', 'purelib', 'platlib']}
 # if in a virtualenv, also exclude the real stdlib location
@@ -146,6 +151,10 @@ def code_filter(self) -> CodeFilter:
         """Default code filter excludes standard library & site-packages."""
         return default_code_filter
 
+    def max_typed_dict_size(self) -> int:
+        """Size up to which a dictionary will be traced as a TypedDict."""
+        return 10
+
 
 def get_default_config() -> Config:
     """Use monkeytype_config.CONFIG if it exists, otherwise DefaultConfig().

monkeytype/encoding.py

@@ -20,11 +20,12 @@
 from monkeytype.db.base import CallTraceThunk
 from monkeytype.exceptions import InvalidTypeError
 from monkeytype.tracing import CallTrace
-from monkeytype.typing import NoneType, NotImplementedType, mappingproxy
+from monkeytype.typing import NoneType, NotImplementedType, is_typed_dict, mappingproxy
 from monkeytype.util import (
     get_func_in_module,
     get_name_in_module,
 )
+from mypy_extensions import TypedDict
 
 
 logger = logging.getLogger(__name__)
@@ -48,13 +49,28 @@
 TypeDict = Dict[str, Any]
 
 
+def typed_dict_to_dict(typ: type) -> TypeDict:
+    elem_types: Dict[str, Any] = {}
+    for k, v in typ.__annotations__.items():
+        elem_types[k] = type_to_dict(v)
+    return {
+        'module': typ.__module__,
+        'qualname': typ.__qualname__,
+        'elem_types': elem_types,
+        'is_typed_dict': True,
+    }
+
+
 def type_to_dict(typ: type) -> TypeDict:
     """Convert a type into a dictionary representation that we can store.
 
     The dictionary must:
         1. Be encodable as JSON
         2. Contain enough information to let us reify the type
     """
+    if is_typed_dict(typ):
+        return typed_dict_to_dict(typ)
+
     # Union and Any are special cases that aren't actually types.
     if is_union(typ):
         qualname = 'Union'
@@ -86,6 +102,12 @@ def type_to_dict(typ: type) -> TypeDict:
 }
 
 
+def typed_dict_from_dict(d: TypeDict) -> type:
+    return TypedDict(d['qualname'],
+                     {k: type_from_dict(v)
+                      for k, v in d['elem_types'].items()})
+
+
 def type_from_dict(d: TypeDict) -> type:
     """Given a dictionary produced by type_to_dict, return the equivalent type.
 
@@ -94,6 +116,8 @@ def type_from_dict(d: TypeDict) -> type:
         InvalidTypeError if the named type isn't actually a type
     """
     module, qualname = d['module'], d['qualname']
+    if d.get('is_typed_dict', False):
+        return typed_dict_from_dict(d)
     if module == 'builtins' and qualname in _HIDDEN_BUILTIN_TYPES:
         typ = _HIDDEN_BUILTIN_TYPES[qualname]
     else: