diff --git a/jedi/evaluate/pep0484.py b/jedi/evaluate/pep0484.py index 80d4b3ca..7b97a89a 100644 --- a/jedi/evaluate/pep0484.py +++ b/jedi/evaluate/pep0484.py @@ -22,7 +22,7 @@ x support for type hint comments for functions, `# type: (int, str) -> int`. import os import re -from parso import ParserSyntaxError +from parso import ParserSyntaxError, parse from parso.python import tree from jedi._compatibility import unicode, force_unicode @@ -80,9 +80,66 @@ def _fix_forward_reference(context, node): return node +def _split_mypy_param_declaration(decl_text): + """ + Split decl_text on commas, but group generic expressions + together. + + For example, given "foo, Bar[baz, biz]" we return + ['foo', 'Bar[baz, biz]']. + + """ + node = parse(decl_text).children[0] + + if node.type == 'name': + return [node.get_code().strip()] + + params = [] + for child in node.children: + if child.type in ['name', 'atom_expr', 'power']: + params.append(child.get_code().strip()) + + return params + + @evaluator_method_cache() def infer_param(execution_context, param): + """ + Infers the type of a function parameter, using type annotations. + """ annotation = param.annotation + if annotation is None: + # If no Python 3-style annotation, look for a Python 2-style comment + # annotation. + # Identify parameters to function in the same sequence as they would + # appear in a type comment. + all_params = [child for child in param.parent.children + if child.type == 'param'] + + node = param.parent.parent + comment = parser_utils.get_following_comment_same_line(node) + if comment is None: + return NO_CONTEXTS + + match = re.match(r"^#\s*type:\s*\(([^#]*)\)\s*->", comment) + if not match: + return NO_CONTEXTS + params_comments = _split_mypy_param_declaration(match.group(1)) + + # Find the specific param being investigated + index = all_params.index(param) + # If the number of parameters doesn't match length of type comment, + # ignore first parameter (assume it's self). + if len(params_comments) != len(all_params): + if index == 0: + # Assume it's self, which is already handled + return NO_CONTEXTS + else: + index -= 1 + param_comment = params_comments[index] + # Construct annotation from type comment + annotation = tree.String(repr(param_comment), node.start_pos) + annotation.parent = node.parent module_context = execution_context.get_root_context() return _evaluate_for_annotation(module_context, annotation) @@ -102,7 +159,26 @@ def py__annotations__(funcdef): @evaluator_method_cache() def infer_return_types(function_context): + """ + Infers the type of a function's return value, + according to type annotations. + """ annotation = py__annotations__(function_context.tree_node).get("return", None) + if annotation is None: + # If there is no Python 3-type annotation, look for a Python 2-type annotation + node = function_context.tree_node + comment = parser_utils.get_following_comment_same_line(node) + if comment is None: + return NO_CONTEXTS + + match = re.match(r"^#\s*type:\s*\([^#]*\)\s*->\s*([^#]*)", comment) + if not match: + return NO_CONTEXTS + + annotation = tree.String( + repr(str(match.group(1).strip())), + node.start_pos) + annotation.parent = node.parent module_context = function_context.get_root_context() return _evaluate_for_annotation(module_context, annotation) diff --git a/jedi/parser_utils.py b/jedi/parser_utils.py index 6b2d6236..990e6cca 100644 --- a/jedi/parser_utils.py +++ b/jedi/parser_utils.py @@ -201,6 +201,9 @@ def get_following_comment_same_line(node): whitespace = node.children[5].get_first_leaf().prefix elif node.type == 'with_stmt': whitespace = node.children[3].get_first_leaf().prefix + elif node.type == 'funcdef': + # actually on the next line + whitespace = node.children[4].get_first_leaf().get_next_leaf().prefix else: whitespace = node.get_last_leaf().get_next_leaf().prefix except AttributeError: diff --git a/test/completion/pep0484_comments.py b/test/completion/pep0484_comments.py index 7707fcc3..87021155 100644 --- a/test/completion/pep0484_comments.py +++ b/test/completion/pep0484_comments.py @@ -47,7 +47,7 @@ b class Employee: pass -from typing import List +from typing import List, Tuple x = [] # type: List[Employee] #? Employee() x[1] @@ -103,3 +103,67 @@ 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 + +# Test instance methods +class Dog: + def __init__(self, age, friends, name): + # type: (int, List[Tuple[str, Dog]], str) -> None + #? int() + self.age = age + self.friends = friends + + #? Dog() + friends[0][1] + + #? str() + self.name = name + + def friend_for_name(self, name): + # type: (str) -> Dog + for (friend_name, friend) in self.friends: + if friend_name == name: + return friend + raise ValueError() + + def bark(self): + pass + +buddy = Dog(UNKNOWN_NAME1, UNKNOWN_NAME2, UNKNOWN_NAME3) +friend = buddy.friend_for_name('buster') +# type of friend is determined by function return type +#! 9 ['def bark'] +friend.bark() + +friend = buddy.friends[0][1] +# type of friend is determined by function parameter type +#! 9 ['def bark'] +friend.bark() + +# type is determined by function parameter type following nested generics +#? str() +friend.name + +# Mypy comment describing function return type. +def annot(): + # type: () -> str + pass + +#? str() +annot() + +# Mypy variable type annotation. +x = UNKNOWN_NAME2 # type: str + +#? str() +x + +class Cat(object): + def __init__(self, age, friends, name): + # type: (int, List[Dog], str) -> None + self.age = age + self.friends = friends + self.name = name + +cat = Cat(UNKNOWN_NAME4, UNKNOWN_NAME5, UNKNOWN_NAME6) +#? str() +cat.name