Codemod for PEP 484 Assign w / type comments -> PEP 526 AnnAssign (#594)

* Codemod for PEP 484 Assign w / type comments -> PEP 526 AnnAssign

Summary:

This codemod is intended to eventually handle all type comments from
PEP 484. This is a partial implementation specifically handling
assignment type comments, which as of PEP 526 are better dealt
with using AnnAssign nodes.

There is more work to do because there are two other kinds of
comments to support: function heading comments and function parameter
inline comments. But the PEP 526 functionality is complete so I feel
like it's worth havign a PR / CI signals / code review at this stage.

Test Plan:

```
python -m unittest libcst.codemod.commands.tests.test_convert_type_comments
```

* Disable on python 3.6, 3.7

The ast module didn't get the `type_comment` information we need
until python 3.8.

It is possible but not a priority right now to enable 3.6 and 3.7
via the typed_ast library, for now I just throw a NotImplementedError
with a nice description. There's a note in the code about where to look
for a typed_ast example in case anyone wants to add support in the
future.

* Fix type errors on the 3.8+ testing fix

* Do a better job of complaining on Python < 3.8

* Updates based on code review

Summary:

Do not strip type comments in the visitor pattern; instead,
reach down from the parent to do it because this makes it
much more reliable that we won't accidentally remove
other comments in a codemod (using visitor state to do this
isn't really feasible once we handle complex statements like
FunctionDef, With, For).

Handle multi-statement statement lines; this works since the
trailing whitespace can only apply to the final statement on
the line. It's not really a critical edge case to handle, but
the code is no more complicated so we might as well.

* Prevent comment stripping for multi-assign

* Note in the docstring that this is a limited WIP

* Reorder checks so the next step will be cleaner

libcst/codemod/commands/convert_type_comments.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import ast
+import builtins
+import functools
+import sys
+from typing import Optional, Set, Union
+
+import libcst as cst
+from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
+
+
+@functools.lru_cache()
+def _empty_module() -> cst.Module:
+    return cst.parse_module("")
+
+
+def _code_for_node(node: cst.CSTNode) -> str:
+    return _empty_module().code_for_node(node)
+
+
+def _ast_for_node(node: cst.CSTNode) -> ast.Module:
+    code = _code_for_node(node)
+    return ast.parse(code, type_comments=True)
+
+
+def _simple_statement_type_comment(
+    node: cst.SimpleStatementLine,
+) -> Optional[str]:
+    return _ast_for_node(node).body[-1].type_comment
+
+
+@functools.lru_cache()
+def _builtins() -> Set[str]:
+    return set(dir(builtins))
+
+
+def _is_builtin(annotation: str) -> bool:
+    return annotation in _builtins()
+
+
+def _convert_annotation(raw: str) -> cst.Annotation:
+    # Convert annotation comments to string annotations to be safe,
+    # otherwise runtime errors would be common.
+    #
+    # Special-case builtins to reduce the amount of quoting noise.
+    #
+    # NOTE: we could potentially detect more cases for skipping quotes
+    # using ScopeProvider, which would make the output prettier.
+    if _is_builtin(raw):
+        return cst.Annotation(annotation=cst.Name(value=raw))
+    else:
+        return cst.Annotation(annotation=cst.SimpleString(f'"{raw}"'))
+
+
+class ConvertTypeComments(VisitorBasedCodemodCommand):
+    """
+    Codemod that converts type comments, as described in
+    https://www.python.org/dev/peps/pep-0484/#type-comments,
+    into PEP 526 annotated assignments.
+
+    This is a work in progress: the codemod only currently handles
+    single-annotation assigns, but it will preserve any type comments
+    that it does not consume.
+    """
+
+    def __init__(self, context: CodemodContext) -> None:
+        if (sys.version_info.major, sys.version_info.minor) < (3, 8):
+            # The ast module did not get `type_comments` until Python 3.7.
+            # In 3.6, we should error than silently running a nonsense codemod.
+            #
+            # NOTE: it is possible to use the typed_ast library for 3.6, but
+            # this is not a high priority right now. See, e.g., the
+            # mypy.fastparse module.
+            raise NotImplementedError(
+                "You are trying to run ConvertTypeComments on a "
+                + "python version without type comment support. Please "
+                + "try using python 3.8+ to run your codemod."
+            )
+        super().__init__(context)
+
+    def _strip_TrailingWhitespace(
+        self,
+        node: cst.TrailingWhitespace,
+    ) -> cst.TrailingWhitespace:
+        return node.with_changes(
+            whitespace=cst.SimpleWhitespace(
+                ""
+            ),  # any whitespace came before the comment, so strip it.
+            comment=None,
+        )
+
+    def _convert_Assign(
+        self,
+        assign: cst.Assign,
+        type_comment: str,
+    ) -> Union[cst.AnnAssign, cst.Assign]:
+        if len(assign.targets) != 1:
+            # this case is not yet implemented, and we short-circuit
+            # it when handling SimpleStatementLine.
+            raise RuntimeError("Should not convert multi-target assign")
+        return cst.AnnAssign(
+            target=assign.targets[0].target,
+            annotation=_convert_annotation(raw=type_comment),
+            value=assign.value,
+            semicolon=assign.semicolon,
+        )
+
+    def leave_SimpleStatementLine(
+        self,
+        original_node: cst.SimpleStatementLine,
+        updated_node: cst.SimpleStatementLine,
+    ) -> cst.SimpleStatementLine:
+        """
+        Convert any SimpleStatementLine containing an Assign with a
+        type comment into a one that uses a PEP 526 AnnAssign.
+        """
+        # determine whether to apply an annotation
+        assign = updated_node.body[-1]
+        if not isinstance(assign, cst.Assign):  # only Assign matters
+            return updated_node
+        type_comment = _simple_statement_type_comment(original_node)
+        if type_comment is None:
+            return updated_node
+        if len(assign.targets) != 1:  # multi-target Assign isn't used
+            return updated_node
+        target = assign.targets[0].target
+        if isinstance(target, cst.Tuple):  # multi-element Assign isn't handled
+            return updated_node
+        # At this point have a single-line Assign with a type comment.
+        # Convert it to an AnnAssign and strip the comment.
+        return updated_node.with_changes(
+            body=[
+                *updated_node.body[:-1],
+                self._convert_Assign(
+                    assign=assign,
+                    type_comment=type_comment,
+                ),
+            ],
+            trailing_whitespace=self._strip_TrailingWhitespace(
+                updated_node.trailing_whitespace
+            ),
+        )

libcst/codemod/commands/tests/test_convert_type_comments.py

@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import sys
+
+from libcst.codemod import CodemodTest
+from libcst.codemod.commands.convert_type_comments import ConvertTypeComments
+
+
+class TestConvertTypeComments(CodemodTest):
+
+    maxDiff = 1000
+    TRANSFORM = ConvertTypeComments
+
+    def assertCodemod38Plus(self, before: str, after: str) -> None:
+        """
+        Assert that the codemod works on Python 3.8+, and that we raise
+        a NotImplementedError on other python versions.
+        """
+        if (sys.version_info.major, sys.version_info.minor) < (3, 8):
+            with self.assertRaises(NotImplementedError):
+                super().assertCodemod(before, after)
+        else:
+            super().assertCodemod(before, after)
+
+    # Tests converting assignment type comments -----------------
+
+    def test_convert_assignments(self) -> None:
+        before = """
+            y = 5  # type: int
+            z = ('this', 7)  # type: typing.Tuple[str, int]
+        """
+        after = """
+            y: int = 5
+            z: "typing.Tuple[str, int]" = ('this', 7)
+        """
+        self.assertCodemod38Plus(before, after)
+
+    def test_convert_assignments_in_context(self) -> None:
+        """
+        Also verify that our matching works regardless of spacing
+        """
+        before = """
+            bar(); baz = 12  # type: int
+
+            def foo():
+                z = ('this', 7) # type: typing.Tuple[str, int]
+
+            class C:
+                attr0 = 10# type: int
+                def __init__(self):
+                    self.attr1 = True  # type: bool
+        """
+        after = """
+            bar(); baz: int = 12
+
+            def foo():
+                z: "typing.Tuple[str, int]" = ('this', 7)
+
+            class C:
+                attr0: int = 10
+                def __init__(self):
+                    self.attr1: bool = True
+        """
+        self.assertCodemod38Plus(before, after)
+
+    def test_no_change_when_type_comment_unused(self) -> None:
+        before = """
+            # type-ignores are not type comments
+            x = 10  # type: ignore
+
+            # a commented type comment (per PEP 484) is not a type comment
+            z = 15  # # type: int
+
+            # a type comment in an illegal location won't be used
+            print("hello")  # type: None
+
+            # We currently cannot handle multiple-target assigns.
+            # Make sure we won't strip those type comments.
+            x, y, z = [], [], []  # type: List[int], List[int], List[str]
+            x, y, z = [], [], []  # type: (List[int], List[int], List[str])
+            a, b, *c = range(5)   # type: float, float, List[float]
+            a, b = 1, 2  # type: Tuple[int, int]
+        """
+        after = before
+        self.assertCodemod38Plus(before, after)