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

* [Pyre] Handle type subscripts when applying annotations.

* Fix lint.

* Use matchers syntax instead of isinstance
This commit is contained in:
Maggie Moss 2020-07-13 17:31:48 -07:00 committed by GitHub
parent 9d3bb11eb8
commit 7219efcd3f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 1 deletions

View file

@ -8,6 +8,7 @@ from dataclasses import dataclass, field
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 @@ class TypeCollector(cst.CSTVisitor):
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 @@ class TypeCollector(cst.CSTVisitor):
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:

View file

@ -624,6 +624,33 @@ class TestApplyAnnotationsVisitor(CodemodTest):
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: