Don't remove annotations on NamedTuples

dflook · dflook · commit f83f45ec6901 · 2020-10-17T10:56:09.000+01:00
diff --git a/src/python_minifier/transforms/remove_annotations.py b/src/python_minifier/transforms/remove_annotations.py
@@ -73,7 +73,25 @@ def is_dataclass_field(node):
 
             return False
 
-        if is_dataclass_field(node):
+        def is_named_tuple(node):
+            if sys.version_info < (3, 5):
+                return False
+
+            if not isinstance(node.parent, ast.ClassDef):
+                return False
+
+            if len(node.parent.bases) == 0:
+                return False
+
+            for node in node.parent.bases:
+                if isinstance(node, ast.Name) and node.id == 'NamedTuple':
+                    return True
+                elif isinstance(node, ast. Attribute) and node.attr == 'NamedTuple':
+                    return True
+
+            return False
+
+        if is_dataclass_field(node) or is_named_tuple(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
@@ -107,3 +107,53 @@ class MyClass:
     expected_ast = ast.parse(expected)
     actual_ast = remove_annotations(source)
     compare_ast(expected_ast, actual_ast)
+
+
+def test_no_remove_namedtuple():
+    if sys.version_info < (3, 6):
+        pytest.skip('annotations unavailable in python < 3.6')
+
+    source = '''
+class MyClass(NamedTuple):
+    myfield: int
+    mysecondfile: str
+
+class MyClass2(typing.NamedTuple):
+    myfield: int
+    mysecondfile: str
+
+class MyClass2(blah.NamedTuple):
+    myfield: int
+    mysecondfile: str
+'''
+    expected = source
+
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_annotations(source)
+    compare_ast(expected_ast, actual_ast)
+
+def test_remove():
+    if sys.version_info < (3, 6):
+        pytest.skip('annotations unavailable in python < 3.6')
+
+    source = '''
+class Dummy(NermedTupel):
+    myfield: int
+    mysecondfile: str
+
+class Dummy(typing.NermedTupel):
+    myfield: int
+    mysecondfile: str
+'''
+    expected = '''
+class Dummy(NermedTupel):
+    myfield: 0
+    mysecondfile: 0
+
+class Dummy(typing.NermedTupel):
+    myfield: 0
+    mysecondfile: 0
+'''
+    expected_ast = ast.parse(expected)
+    actual_ast = remove_annotations(source)
+    compare_ast(expected_ast, actual_ast)
