diff --git a/jedi/api/classes.py b/jedi/api/classes.py index d7dae519..936a70bd 100644 --- a/jedi/api/classes.py +++ b/jedi/api/classes.py @@ -319,7 +319,7 @@ class BaseDefinition(object): call_path = list(stmt_el.generate_call_path()) for i, element in enumerate(call_path): if element is name_part: - return call_path[:i+1] + return call_path[:i + 1] if not isinstance(self._definition, pr.NamePart): raise TypeError('Definition is not a NamePart.') diff --git a/jedi/evaluate/helpers.py b/jedi/evaluate/helpers.py index c4122e69..81200244 100644 --- a/jedi/evaluate/helpers.py +++ b/jedi/evaluate/helpers.py @@ -240,9 +240,22 @@ def statement_elements_in_statement(stmt): stmt_el = stmt_el.next stmt_els = [] - for item in stmt.expression_list(): + for as_name in stmt.as_names: + # TODO This creates a custom pr.Call, we shouldn't do that. + stmt_els.append(pr.Call(as_name._sub_module, as_name, + as_name.start_pos, as_name.end_pos)) + + ass_items = chain.from_iterable(items for items, op in stmt.assignment_details) + for item in stmt.expression_list() + list(ass_items): if isinstance(item, pr.StatementElement): search_stmt_el(item, stmt_els) + elif isinstance(item, pr.ListComprehension): + for stmt in (item.stmt, item.middle, item.input): + stmt_els.extend(statement_elements_in_statement(stmt)) + elif isinstance(item, pr.Lambda): + for stmt in item.params + item.returns: + stmt_els.extend(statement_elements_in_statement(stmt)) + return stmt_els diff --git a/test/test_evaluate/test_helpers.py b/test/test_evaluate/test_helpers.py index 651fe439..0844af20 100644 --- a/test/test_evaluate/test_helpers.py +++ b/test/test_evaluate/test_helpers.py @@ -16,6 +16,12 @@ def test_statement_elements_in_statement(): p = Parser(string) return helpers.statement_elements_in_statement(p.module.statements[0]) + # list comprehension stmt_els = get_stmt_els('foo = [(bar(f), f) for f in baz]') # stmt_els: count all names: 6; + count all arrays: 2 = 8 assert len(stmt_els) == 8 + + # lambda + stmt_els = get_stmt_els('foo = [lambda x: y]') + # stmt_els: count all names: 3; + count all arrays: 1 = 4 + assert len(stmt_els) == 4