From 5cc27f632d5c3d7dd7de08c40e2fb20c7a003baf Mon Sep 17 00:00:00 2001 From: Dave Halter Date: Sat, 2 Jan 2016 21:46:14 +0100 Subject: [PATCH] Improve dict comprehension support. --- jedi/evaluate/iterable.py | 82 +++++++++++++++++++++---------- test/completion/comprehensions.py | 16 ++++-- 2 files changed, 68 insertions(+), 30 deletions(-) diff --git a/jedi/evaluate/iterable.py b/jedi/evaluate/iterable.py index 69594e40..fe8f587e 100644 --- a/jedi/evaluate/iterable.py +++ b/jedi/evaluate/iterable.py @@ -38,7 +38,7 @@ class IterableWrapper(tree.Base): @memoize_default() def _get_names_dict(self, names_dict): try: - builtin_methods = self._builtin_methods + builtin_methods = self.builtin_methods except AttributeError: return names_dict @@ -71,10 +71,10 @@ class BuiltinMethod(IterableWrapper): def has_builtin_methods(cls): - cls._builtin_methods = {} + cls.builtin_methods = {} for func in cls.__dict__.values(): try: - cls._builtin_methods.update(func.registered_builtin_methods) + cls.builtin_methods.update(func.registered_builtin_methods) except AttributeError: pass return cls @@ -166,7 +166,7 @@ class Comprehension(IterableWrapper): return self._get_comprehension().children[1] @memoize_default() - def eval_node(self): + def _eval_node(self, index=0): """ The first part `x + 1` of the list comprehension: @@ -175,9 +175,10 @@ class Comprehension(IterableWrapper): comp_for = self._get_comp_for() # For nested comprehensions we need to search the last one. last_comp = list(comp_for.get_comp_fors())[-1] - return helpers.deep_ast_copy(self._get_comprehension().children[0], parent=last_comp) + return helpers.deep_ast_copy(self._get_comprehension().children[index], parent=last_comp) - def py__iter__(self): + @memoize_default() + def _iterate(self): def nested(comp_fors): comp_for = comp_fors[0] input_node = comp_for.children[3] @@ -192,7 +193,11 @@ class Comprehension(IterableWrapper): for result in nested(comp_fors[1:]): yield result except IndexError: - yield evaluator.eval_element(self.eval_node()) + iterated = evaluator.eval_element(self._eval_node()) + if self.type == 'dict': + yield iterated, evaluator.eval_element(self._eval_node(2)) + else: + yield iterated finally: del evaluator.predefined_if_name_dict_dict[comp_for] @@ -201,10 +206,14 @@ class Comprehension(IterableWrapper): for result in nested(comp_fors): yield result + def py__iter__(self): + return self._iterate() + def __repr__(self): return "<%s of %s>" % (type(self).__name__, self._atom) +@has_builtin_methods class ArrayMixin(object): @memoize_default() def names_dicts(self, search_global=False): # Always False. @@ -213,7 +222,7 @@ class ArrayMixin(object): # builtins only have one class -> [0] scopes = self._evaluator.execute_evaluated(scope, self) names_dicts = list(scopes)[0].names_dicts(search_global) - yield names_dicts[0] + #yield names_dicts[0] yield self._get_names_dict(names_dicts[1]) def py__bool__(self): @@ -230,6 +239,23 @@ class ArrayMixin(object): def name(self): return FakeSequence(self._evaluator, [], self.type).name + @memoize_default() + def dict_values(self): + return unite(self._evaluator.eval_element(v) for k, v in self._items()) + + @register_builtin_method('values', type='dict') + def _imitate_values(self): + items = self.dict_values() + return create_evaluated_sequence_set(self._evaluator, items, type='list') + #return set([FakeSequence(self._evaluator, [AlreadyEvaluated(items)], 'tuple')]) + + @register_builtin_method('items', type='dict') + def _imitate_items(self): + items = [set([FakeSequence(self._evaluator, (k, v), 'tuple')]) + for k, v in self._items()] + + return create_evaluated_sequence_set(self._evaluator, *items, type='list') + class ListComprehension(Comprehension, ArrayMixin): type = 'list' @@ -243,18 +269,39 @@ class SetComprehension(Comprehension, ArrayMixin): type = 'set' +@has_builtin_methods class DictComprehension(Comprehension, ArrayMixin): type = 'dict' def _get_comp_for(self): return self._get_comprehension().children[3] + def py__getitem__(self, index): + for keys, values in self._iterate(): + for k in keys: + if isinstance(k, compiled.CompiledObject): + if k.obj == index: + return values + return self.dict_values() + + def dict_values(self): + return unite(values for keys, values in self._iterate()) + + @register_builtin_method('items', type='dict') + def _imitate_items(self): + items = set(FakeSequence(self._evaluator, + (AlreadyEvaluated(keys), AlreadyEvaluated(values)), 'tuple') + for keys, values in self._iterate()) + + return create_evaluated_sequence_set(self._evaluator, items, type='list') + + + class GeneratorComprehension(Comprehension, GeneratorMixin): pass -@has_builtin_methods class Array(IterableWrapper, ArrayMixin): mapping = {'(': 'tuple', '[': 'list', @@ -277,23 +324,6 @@ class Array(IterableWrapper, ArrayMixin): def name(self): return helpers.FakeName(self.type, parent=self) - @memoize_default() - def dict_values(self): - return unite(self._evaluator.eval_element(v) for k, v in self._items()) - - @register_builtin_method('values', type='dict') - def _imitate_values(self): - items = unite(self._evaluator.eval_element(v) for k, v in self._items()) - return create_evaluated_sequence_set(self._evaluator, items, type='list') - #return set([FakeSequence(self._evaluator, [AlreadyEvaluated(items)], 'tuple')]) - - @register_builtin_method('items', type='dict') - def _imitate_items(self): - items = set(FakeSequence(self._evaluator, (k, v), 'tuple') - for k, v in self._items()) - - return create_evaluated_sequence_set(self._evaluator, items, type='list') - def py__getitem__(self, index): """Here the index is an int/str. Raises IndexError/KeyError.""" if self.type == 'dict': diff --git a/test/completion/comprehensions.py b/test/completion/comprehensions.py index b1e0caf2..83dba5da 100644 --- a/test/completion/comprehensions.py +++ b/test/completion/comprehensions.py @@ -126,13 +126,21 @@ right # dict comprehensions # ----------------- -d = {a - 1: b for a, b in {1: 'a', 3: 1.0}.items()} -#? str() -list()[0] - #? int() list({a - 1: 3 for a in [1]})[0] +d = {a - 1: b for a, b in {1: 'a', 3: 1.0}.items()} +#? int() +list(d)[0] +#? str() +d.values()[0] +#? str() +d[0] +#? float() str() +d[1] +#? float() +d[2] + # ----------------- # set comprehensions # -----------------