bpo-38870: Implement round tripping support for typed AST in ast.unparse (GH-17797)

isidentical · web-flow · commit dff92bb31f7d · 2020-05-17T00:04:12.000+01:00
diff --git a/Lib/ast.py b/Lib/ast.py
@@ -648,6 +648,7 @@ def __init__(self):
         self._source = []
         self._buffer = []
         self._precedences = {}
+        self._type_ignores = {}
         self._indent = 0
 
     def interleave(self, inter, f, seq):
@@ -697,11 +698,15 @@ def buffer(self):
         return value
 
     @contextmanager
-    def block(self):
+    def block(self, *, extra = None):
         """A context manager for preparing the source for blocks. It adds
         the character':', increases the indentation on enter and decreases
-        the indentation on exit."""
+        the indentation on exit. If *extra* is given, it will be directly
+        appended after the colon character.
+        """
         self.write(":")
+        if extra:
+            self.write(extra)
         self._indent += 1
         yield
         self._indent -= 1
@@ -748,6 +753,11 @@ def get_raw_docstring(self, node):
         if isinstance(node, Constant) and isinstance(node.value, str):
             return node
 
+    def get_type_comment(self, node):
+        comment = self._type_ignores.get(node.lineno) or node.type_comment
+        if comment is not None:
+            return f" # type: {comment}"
+
     def traverse(self, node):
         if isinstance(node, list):
             for item in node:
@@ -770,7 +780,12 @@ def _write_docstring_and_traverse_body(self, node):
             self.traverse(node.body)
 
     def visit_Module(self, node):
+        self._type_ignores = {
+            ignore.lineno: f"ignore{ignore.tag}"
+            for ignore in node.type_ignores
+        }
         self._write_docstring_and_traverse_body(node)
+        self._type_ignores.clear()
 
     def visit_FunctionType(self, node):
         with self.delimit("(", ")"):
@@ -811,6 +826,8 @@ def visit_Assign(self, node):
             self.traverse(target)
             self.write(" = ")
         self.traverse(node.value)
+        if type_comment := self.get_type_comment(node):
+            self.write(type_comment)
 
     def visit_AugAssign(self, node):
         self.fill()
@@ -966,7 +983,7 @@ def _function_helper(self, node, fill_suffix):
         if node.returns:
             self.write(" -> ")
             self.traverse(node.returns)
-        with self.block():
+        with self.block(extra=self.get_type_comment(node)):
             self._write_docstring_and_traverse_body(node)
 
     def visit_For(self, node):
@@ -980,7 +997,7 @@ def _for_helper(self, fill, node):
         self.traverse(node.target)
         self.write(" in ")
         self.traverse(node.iter)
-        with self.block():
+        with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
         if node.orelse:
             self.fill("else")
@@ -1018,13 +1035,13 @@ def visit_While(self, node):
     def visit_With(self, node):
         self.fill("with ")
         self.interleave(lambda: self.write(", "), self.traverse, node.items)
-        with self.block():
+        with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
 
     def visit_AsyncWith(self, node):
         self.fill("async with ")
         self.interleave(lambda: self.write(", "), self.traverse, node.items)
-        with self.block():
+        with self.block(extra=self.get_type_comment(node)):
             self.traverse(node.body)
 
     def visit_JoinedStr(self, node):
diff --git a/Lib/test/test_unparse.py b/Lib/test/test_unparse.py
@@ -108,12 +108,12 @@ class Foo: pass
     suite1
 """
 
-docstring_prefixes = [
+docstring_prefixes = (
     "",
     "class foo:\n    ",
     "def foo():\n    ",
     "async def foo():\n    ",
-]
+)
 
 class ASTTestCase(unittest.TestCase):
     def assertASTEqual(self, ast1, ast2):
@@ -340,6 +340,37 @@ def test_function_type(self):
         ):
             self.check_ast_roundtrip(function_type, mode="func_type")
 
+    def test_type_comments(self):
+        for statement in (
+            "a = 5 # type:",
+            "a = 5 # type: int",
+            "a = 5 # type: int and more",
+            "def x(): # type: () -> None\n\tpass",
+            "def x(y): # type: (int) -> None and more\n\tpass",
+            "async def x(): # type: () -> None\n\tpass",
+            "async def x(y): # type: (int) -> None and more\n\tpass",
+            "for x in y: # type: int\n\tpass",
+            "async for x in y: # type: int\n\tpass",
+            "with x(): # type: int\n\tpass",
+            "async with x(): # type: int\n\tpass"
+        ):
+            self.check_ast_roundtrip(statement, type_comments=True)
+
+    def test_type_ignore(self):
+        for statement in (
+            "a = 5 # type: ignore",
+            "a = 5 # type: ignore and more",
+            "def x(): # type: ignore\n\tpass",
+            "def x(y): # type: ignore and more\n\tpass",
+            "async def x(): # type: ignore\n\tpass",
+            "async def x(y): # type: ignore and more\n\tpass",
+            "for x in y: # type: ignore\n\tpass",
+            "async for x in y: # type: ignore\n\tpass",
+            "with x(): # type: ignore\n\tpass",
+            "async with x(): # type: ignore\n\tpass"
+        ):
+            self.check_ast_roundtrip(statement, type_comments=True)
+
 
 class CosmeticTestCase(ASTTestCase):
     """Test if there are cosmetic issues caused by unnecesary additions"""
