def test_quoted_ast_annassign(self): example_code = "a: 'DjangoModel'" tree = ast.parse(example_code) objects = get_ast_objects(tree) assert type(objects[0]) == ast.AnnAssign annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations[0]['annotation'] == 'DjangoModel'
def test_ast_name(self): example_code = 'def test() -> DjangoModel:\n' ' pass' tree = ast.parse(example_code) objects = get_ast_objects(tree) assert type(objects[0]) == ast.Name annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations[0]['annotation'] == 'DjangoModel'
def test_star_arg(self): example_code = textwrap.dedent(''' from typing import List x: List def b(*, x: List[str]): pass ''') tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations == [ { 'annotation': 'List', 'line_number': 4 }, { 'annotation': 'str', 'line_number': 7 }, { 'annotation': 'List', 'line_number': 7 }, ]
def test_custom_type(self): example_code = 'x = Tuple[float, float, float]' tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = flatten_list([get_annotations(obj) for obj in objects]) just_annotations = [i['annotation'] for i in annotations] assert just_annotations == []
def test_simple_typing_import(self): example_code = 'import typing\n\nx: typing.List' tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = flatten_list([get_annotations(obj) for obj in objects]) just_annotations = [i['annotation'] for i in annotations] assert just_annotations == ['typing.List']
def test_ast_subscript(self): example_code = "def test() -> Literal['one', 'two', 'three']:\n" ' pass' tree = ast.parse(example_code) objects = get_ast_objects(tree) assert type(objects[0]) == ast.Subscript annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations[0]['annotation'] == 'one' assert annotations[1]['annotation'] == 'two' assert annotations[2]['annotation'] == 'three' assert annotations[3]['annotation'] == 'Literal'
def test_get_complex_variable_annotation(self): example_code = 'a: Dict[Dict[Pattern, MutableSet], Tuple[Optional[str], Sequence]] = {}' tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = flatten_list([get_annotations(obj) for obj in objects]) just_annotations = [i['annotation'] for i in annotations] assert just_annotations.count('Dict') == 2 assert just_annotations.count('Tuple') == 1 assert just_annotations.count('str') == 1 assert just_annotations.count('Sequence') == 1 assert just_annotations.count('Pattern') == 1 assert just_annotations.count('MutableSet') == 1
def test_ast_attribute(self): example_code = 'variable: typing.Dict[typing.Dict[typing.Tuple[dict, list], str], int]' tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = flatten_list([get_annotations(obj) for obj in objects]) just_annotations = [i['annotation'] for i in annotations] assert just_annotations.count('typing.Dict') == 2 assert just_annotations.count('typing.Tuple') == 1 assert just_annotations.count('dict') == 1 assert just_annotations.count('list') == 1 assert just_annotations.count('str') == 1 assert just_annotations.count('int') == 1
def test_deep_attribute(self): """ An ast.Attribute with more than 2 levels should return an empty dict. """ example_code = textwrap.dedent(''' import pandas as pd def b(x: pd.core.frame.DataFrame): pass ''') tree = ast.parse(example_code) objects = get_ast_objects(tree) annotations = get_annotations(objects[0]) assert annotations == {}
def test_ast_tuple(self): example_code = "def test() -> Literal['one', 'two', 'three']:\n" ' pass' tree = ast.parse(example_code) objects = get_ast_objects(tree) tuple_object = objects[0].slice if sys.version_info < (3, 9): assert type(tuple_object) == _ast.Index else: assert type(tuple_object) == ast.Tuple annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations[0]['annotation'] == 'one' assert annotations[1]['annotation'] == 'two' assert annotations[2]['annotation'] == 'three' assert annotations[3]['annotation'] == 'Literal' assert len(annotations) == 4
def test_get_ast_objects(self): example_code = textwrap.dedent(""" a: dict = {} def test(x: int, y: str) -> None: pass """) """ From the example code we've set up above, we expect to receive 4 useful objects: - The `a: dict` annotation assignment (ast.AnnAssign type) - The `x: int` argument (ast.Name type) - The `y: str` argument (ast.Name type) - The `-> None` return value (ast.Constant type) If we don't, our hook wouldn't work. We also expect them in this order. """ tree = ast.parse(example_code) objects = get_ast_objects(tree) assert type( objects[0]) == ast.AnnAssign and objects[0].annotation.id == 'dict' assert type(objects[1]) == ast.Name and objects[1].id == 'int' assert type(objects[2]) == ast.Name and objects[2].id == 'str' assert type(objects[3]) == ast.Constant and objects[3].value is None """ Since we know the annotations we're expecting here, we can also quickly add another sub-test. """ annotations = flatten_list([get_annotations(obj) for obj in objects]) assert annotations[0]['annotation'] == 'dict' assert annotations[1]['annotation'] == 'int' assert annotations[2]['annotation'] == 'str' assert len(annotations) == 3
def test_bad_ast_type(self): assert get_annotations('test') is None