Don't remove annotations from TypedDict & NamedTuple

dflook · dflook · commit 9100a93543a3 · 2020-10-17T10:56:09.000+01:00
diff --git a/docs/source/transforms/remove_annotations.rst b/docs/source/transforms/remove_annotations.rst
@@ -3,8 +3,16 @@ Remove Annotations
 
 This transform removes function annotations and variable annotations.
 
-This transform is generally safe to use. Although the annotations have no meaning to python,
-they are made available at runtime. If you know the module requires the annotations to be kept, disable this transform.
+This transform is generally safe to use. Although the annotations have no meaning to the python language,
+they are made available at runtime. Some python library features require annotations to be kept.
+
+If these are detected, annotations are kept for that class:
+
+    - dataclasses.dataclass
+    - typing.NamedTuple
+    - typing.NamedDict
+
+If you know the module requires the annotations to be kept, disable this transform.
 
 If a variable annotation without assignment is used the annotation is changed to a literal zero instead of being removed.
 
diff --git a/src/python_minifier/transforms/remove_annotations.py b/src/python_minifier/transforms/remove_annotations.py
@@ -73,7 +73,7 @@ def is_dataclass_field(node):
 
             return False
 
-        def is_named_tuple(node):
+        def is_typing_sensitive(node):
             if sys.version_info < (3, 5):
                 return False
 
@@ -83,15 +83,17 @@ def is_named_tuple(node):
             if len(node.parent.bases) == 0:
                 return False
 
+            tricky_types = ['NamedTuple', 'TypedDict']
+
             for node in node.parent.bases:
-                if isinstance(node, ast.Name) and node.id == 'NamedTuple':
+                if isinstance(node, ast.Name) and node.id in tricky_types:
                     return True
-                elif isinstance(node, ast. Attribute) and node.attr == 'NamedTuple':
+                elif isinstance(node, ast. Attribute) and node.attr in tricky_types:
                     return True
 
             return False
 
-        if is_dataclass_field(node) or is_named_tuple(node):
+        if is_dataclass_field(node) or is_typing_sensitive(node):
             return node
         elif node.value:
             return self.add_child(ast.Assign([node.target], node.value), parent=node.parent)
diff --git a/test/test_remove_annotations.py b/test/test_remove_annotations.py
@@ -157,3 +157,47 @@ class Dummy(typing.NermedTupel):
     expected_ast = ast.parse(expected)
     actual_ast = remove_annotations(source)
     compare_ast(expected_ast, actual_ast)
+
+
+def test_no_remove_typeddict():
+    if sys.version_info < (3, 8):
+        pytest.skip('annotations unavailable in python < 3.6')
+
+    source = '''
+class Dummy(TypedDict):
+    myfield: int
+    mysecondfile: str
+
+class Dummy(HypedDict):
+    myfield: int
+    mysecondfile: str
+
+class Dummy(typing.TypedDict):
+    myfield: int
+    mysecondfile: str
+
+class Dummy(typing.TypedDic):
+    myfield: int
+    mysecondfile: str
+'''
+    expected = '''
+class Dummy(TypedDict):
+    myfield: int
+    mysecondfile: str
+    
+class Dummy(HypedDict):
+    myfield: 0
+    mysecondfile: 0    
+
+class Dummy(typing.TypedDict):
+    myfield: int
+    mysecondfile: str
+    
+class Dummy(typing.TypedDic):
+    myfield: 0
+    mysecondfile: 0
+        
+'''
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_annotations(source)
+    compare_ast(expected_ast, actual_ast)
