[Pyre] Handle type subscripts when applying annotations. (Instagram#335)

* [Pyre] Handle type subscripts when applying annotations.

* Fix lint.

* Use matchers syntax instead of isinstance

libcst/codemod/visitors/_apply_type_annotations.py

@@ -8,6 +8,7 @@
 from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
 
 import libcst as cst
+from libcst import matchers as m
 from libcst.codemod._context import CodemodContext
 from libcst.codemod._visitor import ContextAwareTransformer
 from libcst.codemod.visitors._add_imports import AddImportsVisitor
@@ -133,6 +134,8 @@ def _handle_Index(self, slice: cst.Index, node: cst.Subscript) -> cst.Subscript:
 
     def _handle_Subscript(self, node: cst.Subscript) -> cst.Subscript:
         slice = node.slice
+        if m.matches(node.value, m.Name(value="Type")):
+            return node
         if isinstance(slice, list):
             new_slice = []
             for item in slice:
@@ -163,7 +166,7 @@ def _create_import_from_annotation(self, returns: cst.Annotation) -> cst.Annotat
             return cst.Annotation(annotation=attr)
         if isinstance(annotation, cst.Subscript):
             value = annotation.value
-            if isinstance(value, cst.Name) and value.value == "Type":
+            if m.matches(value, m.Name(value="Type")):
                 return returns
             return cst.Annotation(annotation=self._handle_Subscript(annotation))
         else:

libcst/codemod/visitors/tests/test_apply_type_annotations.py

@@ -624,6 +624,33 @@ def foo() -> typing.Sequence[int]:
                   return []
                 """,
             ),
+            (
+                """
+                from typing import Dict
+
+                example: Dict[str, Type[foo.Example]] = ...
+                """,
+                """
+                from typing import Type
+
+                def foo() -> Type[foo.Example]:
+                    class Example:
+                        pass
+                    return Example
+
+                example = { "test": foo() }
+                """,
+                """
+                from typing import Dict, Type
+
+                def foo() -> Type[foo.Example]:
+                    class Example:
+                        pass
+                    return Example
+
+                example: Dict[str, Type[foo.Example]] = { "test": foo() }
+                """,
+            ),
         )
     )
     def test_annotate_functions(self, stub: str, before: str, after: str) -> None: