diff --git a/jedi/evaluate/__init__.py b/jedi/evaluate/__init__.py index b307eddc..82c25412 100644 --- a/jedi/evaluate/__init__.py +++ b/jedi/evaluate/__init__.py @@ -211,6 +211,13 @@ class Evaluator(object): if c[0] == '(' and (not pr.is_node(c[1], 'testlist_comp') or c[1].children[1] != ','): return self.eval_element(c[1]) + try: + comp_for = c[1].children[1] + except (IndexError, AttributeError): + pass + else: + if isinstance(comp_for, pr.CompFor): + return [iterable.Comprehension(self, atom)] return [iterable.Array(self, atom)] def _eval_trailer(self, types, trailer): diff --git a/jedi/evaluate/finder.py b/jedi/evaluate/finder.py index 58cc599f..8b94511f 100644 --- a/jedi/evaluate/finder.py +++ b/jedi/evaluate/finder.py @@ -74,6 +74,8 @@ class NameFinder(object): except KeyError: return [] + if isinstance(scope, pr.CompFor): + return names names = pr.filter_after_position(names, position) names = [name for name in names if name.is_definition()] @@ -285,9 +287,13 @@ class NameFinder(object): for_types = self._evaluator.eval_element(typ.children[-3]) for_types = iterable.get_iterator_types(for_types) types += check_tuple_assignments(for_types, name) + elif typ.isinstance(pr.CompFor): + for_types = self._evaluator.eval_element(typ.children[3]) + for_types = iterable.get_iterator_types(for_types) + types += check_tuple_assignments(for_types, name) elif isinstance(typ, pr.Param): types += self._eval_param(typ) - elif typ.isinstance(pr.ExprStmt, pr.CompFor): + elif typ.isinstance(pr.ExprStmt): types += self._remove_statements(typ, name) elif isinstance(typ, pr.Import): types += imports.ImportWrapper(self._evaluator, name).follow() diff --git a/jedi/evaluate/iterable.py b/jedi/evaluate/iterable.py index a9753bb5..0f469cbd 100644 --- a/jedi/evaluate/iterable.py +++ b/jedi/evaluate/iterable.py @@ -119,6 +119,29 @@ class GeneratorComprehension(Generator): return self._evaluator.eval_statement_element(self.comprehension) +class Comprehension(IterableWrapper): + def __init__(self, evaluator, atom): + self._evaluator = evaluator + self._atom = atom + + def get_index_types(self, evaluator, index): + return self._evaluator.eval_element(self.eval_node()) + + @memoize_default() + def eval_node(self): + """ + The first part `x + 1` of the list comprehension: + + [x + 1 for x in foo] + """ + comprehension = self._atom.children[1] + c = comprehension.children + return helpers.deep_ast_copy(c[0], {comprehension: c[1]}) + + def __repr__(self): + return "" % (type(self).__name__, self._atom) + + class Array(IterableWrapper): """ Used as a mirror to pr.Array, if needed. It defines some getter diff --git a/test/completion/basic.py b/test/completion/basic.py index 7bba8123..0a6e45e5 100644 --- a/test/completion/basic.py +++ b/test/completion/basic.py @@ -126,7 +126,7 @@ a = [a for a in [1]] #? int() a[0] -a = [a for a in 1,2] +a = [a for a in (1, 2)] #? int() a[0]