Preserve type annotations by converting them to comments

py_backwards/compiler.py

@@ -3,7 +3,7 @@
 from traceback import format_exc
 from typing import List, Tuple, Optional
 from typed_ast import ast3 as ast
-from astunparse import unparse, dump
+from typed_astunparse import unparse, dump
 from autopep8 import fix_code
 from .files import get_input_output_paths, InputOutput
 from .transformers import transformers

py_backwards/transformers/functions_annotations.py

@@ -1,4 +1,5 @@
 from typed_ast import ast3 as ast
+from typed_astunparse import unparse
 from .base import BaseNodeTransformer
 
 
@@ -8,17 +9,34 @@ def fn(x: int) -> int:
             pass
     To:
         def fn(x):
+            # type: (int) -> int
             pass
 
     """
     target = (2, 7)
 
     def visit_arg(self, node: ast.arg) -> ast.arg:
         self._tree_changed = True
+        node._annotation = node.annotation
         node.annotation = None
         return self.generic_visit(node)  # type: ignore
 
     def visit_FunctionDef(self, node: ast.FunctionDef):
         self._tree_changed = True
+
+        returns = node.returns
         node.returns = None
-        return self.generic_visit(node)  # type: ignore
+
+        node = self.generic_visit(node)  # type: ignore
+
+        args_types = [unparse(arg._annotation).strip() for arg in node.args.args if arg._annotation is not None]
+        if getattr(node.args.vararg, '_annotation', None) is not None:
+            args_types += ['*' + unparse(node.args.vararg._annotation).strip()]
+        args_types += [unparse(arg._annotation).strip() for arg in node.args.kwonlyargs if arg._annotation is not None]
+        if getattr(node.args.kwarg, '_annotation', None) is not None:
+            args_types += ['**' + unparse(node.args.kwarg._annotation).strip()]
+
+        if returns:
+            node.type_comment = '(' + (', '.join(args_types)) + ') -> ' + unparse(returns).strip()
+
+        return node

py_backwards/transformers/variables_annotations.py

@@ -11,7 +11,7 @@ class VariablesAnnotationsTransformer(BaseTransformer):
         a: int = 10
         b: int
     To:
-        a = 10
+        a = 10  # type: int
 
     """
     target = (3, 5)

setup.py

@@ -3,7 +3,7 @@
 
 VERSION = '0.7'
 
-install_requires = ['typed-ast', 'autopep8', 'colorama', 'py-backwards-astunparse']
+install_requires = ['typed-ast', 'autopep8', 'colorama', 'typed_astunparse']
 extras_require = {':python_version<"3.4"': ['pathlib2'],
                   ':python_version<"3.5"': ['typing']}
 

tests/transformers/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 from types import ModuleType
 from typed_ast.ast3 import parse, dump
-from astunparse import unparse, dump as dump_pretty
+from typed_astunparse import unparse, dump as dump_pretty
 
 
 @pytest.fixture

tests/utils/test_snippet.py

@@ -1,5 +1,5 @@
 from typed_ast import ast3 as ast
-from astunparse import unparse
+from typed_astunparse import unparse
 from py_backwards.utils.snippet import (snippet, let, find_variables,
                                         VariablesReplacer, extend_tree)
 

tests/utils/test_tree.py

@@ -1,5 +1,5 @@
 from typed_ast import ast3 as ast
-from astunparse import unparse
+from typed_astunparse import unparse
 from py_backwards.utils.snippet import snippet
 from py_backwards.utils.tree import (get_parent, get_node_position,
                                      find, insert_at, replace_at)