diff --git a/jedi/evaluate/finder.py b/jedi/evaluate/finder.py index 5003848c..4114518a 100644 --- a/jedi/evaluate/finder.py +++ b/jedi/evaluate/finder.py @@ -304,6 +304,14 @@ class NameFinder(object): @memoize_default(set(), evaluator_is_first_arg=True) def _name_to_types(evaluator, name, scope): typ = name.get_definition() + if typ.isinstance(tree.ForStmt): + types = pep0484.find_type_from_comment_hint_for(evaluator, typ, name) + if types: + return types + if typ.isinstance(tree.WithStmt): + types = pep0484.find_type_from_comment_hint_with(evaluator, typ, name) + if types: + return types if typ.isinstance(tree.ForStmt, tree.CompFor): container_types = evaluator.eval_element(typ.children[3]) for_types = iterable.py__iter__types(evaluator, container_types, typ.children[3]) @@ -355,6 +363,10 @@ def _remove_statements(evaluator, stmt, name): check_instance = stmt.instance stmt = stmt.var + pep0484types = \ + pep0484.find_type_from_comment_hint_assign(evaluator, stmt, name) + if pep0484types: + return pep0484types types |= evaluator.eval_statement(stmt, seek_name=name) if check_instance is not None: diff --git a/jedi/evaluate/pep0484.py b/jedi/evaluate/pep0484.py index 08588d71..dc951aaf 100644 --- a/jedi/evaluate/pep0484.py +++ b/jedi/evaluate/pep0484.py @@ -9,13 +9,14 @@ v Function parameter annotations with builtin/custom type classes v Function returntype annotations with builtin/custom type classes v Function parameter annotations with strings (forward reference) v Function return type annotations with strings (forward reference) -x Local variable type hints +v Local variable type hints v Assigned types: `Url = str\ndef get(url:Url) -> str:` -x Type hints in `with` statements +v Type hints in `with` statements x Stub files support x support `@no_type_check` and `@no_type_check_decorator` -x support for type hint comments `# type: (int, str) -> int`. See comment from - Guido https://github.com/davidhalter/jedi/issues/662 +x support for typing.cast() operator +x support for type hint comments for functions, `# type: (int, str) -> int`. + See comment from Guido https://github.com/davidhalter/jedi/issues/662 """ import itertools @@ -28,12 +29,23 @@ from jedi.common import unite from jedi.evaluate import compiled from jedi import debug from jedi import _compatibility +import re -def _evaluate_for_annotation(evaluator, annotation): +def _evaluate_for_annotation(evaluator, annotation, index=None): + """ + Evaluates a string-node, looking for an annotation + If index is not None, the annotation is expected to be a tuple + and we're interested in that index + """ if annotation is not None: definitions = evaluator.eval_element( _fix_forward_reference(evaluator, annotation)) + if index is not None: + definitions = list(itertools.chain.from_iterable( + definition.py__getitem__(index) for definition in definitions + if definition.type == 'tuple' and + len(list(definition.py__iter__())) >= index)) return list(itertools.chain.from_iterable( evaluator.execute(d) for d in definitions)) else: @@ -136,3 +148,48 @@ def get_types_for_typing_module(evaluator, typ, node): result = evaluator.execute_evaluated(factory, compiled_classname, args) return result + + +def find_type_from_comment_hint_for(evaluator, node, name): + return \ + _find_type_from_comment_hint(evaluator, node, node.children[1], name) + + +def find_type_from_comment_hint_with(evaluator, node, name): + assert len(node.children[1].children) == 3, \ + "Can only be here when children[1] is 'foo() as f'" + return _find_type_from_comment_hint( + evaluator, node, node.children[1].children[2], name) + + +def find_type_from_comment_hint_assign(evaluator, node, name): + return \ + _find_type_from_comment_hint(evaluator, node, node.children[0], name) + + +def _find_type_from_comment_hint(evaluator, node, varlist, name): + index = None + if varlist.type in ("testlist_star_expr", "exprlist"): + # something like "a, b = 1, 2" + index = 0 + for child in varlist.children: + if child == name: + break + if child.type == "operator": + continue + index += 1 + else: + return [] + + comment = node.get_following_comment_same_line() + if comment is None: + return [] + match = re.match(r"^#\s*type:\s*([^#]*)", comment) + if not match: + return [] + annotation = tree.String( + tree.zero_position_modifier, + repr(str(match.group(1).strip())), + node.start_pos) + annotation.parent = node.parent + return _evaluate_for_annotation(evaluator, annotation, index) diff --git a/jedi/parser/tree.py b/jedi/parser/tree.py index 89c37841..133b75ae 100644 --- a/jedi/parser/tree.py +++ b/jedi/parser/tree.py @@ -239,13 +239,35 @@ class Leaf(Base): else: node = c[i - 1] break - while True: try: node = node.children[-1] except AttributeError: # A Leaf doesn't have children. return node + def get_next(self): + """ + Returns the next leaf in the parser tree. + """ + node = self + while True: + c = node.parent.children + i = c.index(node) + try: + node = c[i + 1] + except IndexError: + node = node.parent + if node.parent is None: + raise IndexError('Cannot access the next element of the last one.') + else: + break + while True: + try: + node = node.children[0] + except AttributeError: # A Leaf doesn't have children. + return node + + def get_code(self, normalized=False): if normalized: return self.value @@ -264,6 +286,7 @@ class Leaf(Base): except IndexError: return None + def prev_sibling(self): """ The node/leaf immediately preceding the invocant in their parent's @@ -277,6 +300,7 @@ class Leaf(Base): return None return self.parent.children[i - 1] + def nodes_to_execute(self, last_added=False): return [] @@ -488,6 +512,39 @@ class BaseNode(Base): except AttributeError: return self.children[0] + def last_leaf(self): + try: + return self.children[-1].last_leaf() + except AttributeError: + return self.children[-1] + + def get_following_comment_same_line(self): + """ + returns (as string) any comment that appears on the same line, + after the node, including the # + """ + try: + if self.isinstance(ForStmt): + whitespace = self.children[5].first_leaf().prefix + elif self.isinstance(WithStmt): + whitespace = self.children[3].first_leaf().prefix + else: + whitespace = self.last_leaf().get_next().prefix + except AttributeError: + return None + except ValueError: + # in some particular cases, the tree doesn't seem to be linked + # correctly + return None + if "#" not in whitespace: + return None + comment = whitespace[whitespace.index("#"):] + if "\r" in comment: + comment = comment[:comment.index("\r")] + if "\n" in comment: + comment = comment[:comment.index("\n")] + return comment + @utf8_repr def __repr__(self): code = self.get_code().replace('\n', ' ').strip() diff --git a/test/completion/pep0484_comments.py b/test/completion/pep0484_comments.py new file mode 100644 index 00000000..7707fcc3 --- /dev/null +++ b/test/completion/pep0484_comments.py @@ -0,0 +1,105 @@ +a = 3 # type: str +#? str() +a + +b = 3 # type: str but I write more +#? int() +b + +c = 3 # type: str # I comment more +#? str() +c + +d = "It should not read comments from the next line" +# type: int +#? str() +d + +# type: int +e = "It should not read comments from the previous line" +#? str() +e + +class BB: pass + +def test(a, b): + a = a # type: BB + c = a # type: str + d = a + # type: str + e = a # type: str # Should ignore long whitespace + + #? BB() + a + #? str() + c + #? BB() + d + #? str() + e + +a,b = 1, 2 # type: str, float +#? str() +a +#? float() +b + +class Employee: + pass + +from typing import List +x = [] # type: List[Employee] +#? Employee() +x[1] +x, y, z = [], [], [] # type: List[int], List[int], List[str] +#? int() +y[2] +x, y, z = [], [], [] # type: (List[float], List[float], List[BB]) +for zi in z: + #? BB() + zi + +x = [ + 1, + 2, +] # type: List[str] + +#? str() +x[1] + + +for bar in foo(): # type: str + #? str() + bar + +for bar, baz in foo(): # type: int, float + #? int() + bar + #? float() + baz + +for bar, baz in foo(): + # type: str, str + """ type hinting on next line should not work """ + #? + bar + #? + baz + +with foo(): # type: int + ... + +with foo() as f: # type: str + #? str() + f + +with foo() as f: + # type: str + """ type hinting on next line should not work """ + #? + f + +aaa = some_extremely_long_function_name_that_doesnt_leave_room_for_hints() \ + # type: float # We should be able to put hints on the next line with a \ +#? float() +aaa