python · pablogsal · May 16, 2020 · Jan 2, 2020 · Mar 17, 2020
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):
@@ -333,6 +333,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"""