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

1 files changed, 33 insertions, 2 deletions

diff --git a/Lib/test/test_unparse.py b/Lib/test/test_unparse.py
index 4f57428..d543ca2 100644
--- a/Lib/test/test_unparse.py
+++ b/Lib/test/test_unparse.py
@@ -108,12 +108,12 @@ with f() as x, g() as y:
     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 @@ class UnparseTestCase(ASTTestCase):
         ):
             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"""