diff --git a/jedi/common/utils.py b/jedi/common/utils.py index 72726a46..92ff52e2 100644 --- a/jedi/common/utils.py +++ b/jedi/common/utils.py @@ -1,4 +1,5 @@ import os +from contextlib import contextmanager def traverse_parents(path, include_current=False): @@ -10,3 +11,16 @@ def traverse_parents(path, include_current=False): yield path previous = path path = os.path.dirname(path) + + +@contextmanager +def monkeypatch(obj, attribute_name, new_value): + """ + Like pytest's monkeypatch, but as a context manager. + """ + old_value = getattr(obj, attribute_name) + try: + setattr(obj, attribute_name, new_value) + yield + finally: + setattr(obj, attribute_name, old_value) diff --git a/jedi/evaluate/lazy_context.py b/jedi/evaluate/lazy_context.py index 9380212e..97e4b068 100644 --- a/jedi/evaluate/lazy_context.py +++ b/jedi/evaluate/lazy_context.py @@ -1,4 +1,5 @@ from jedi.evaluate.base_context import ContextSet, NO_CONTEXTS +from jedi.common.utils import monkeypatch class AbstractLazyContext(object): def __init__(self, data): @@ -40,12 +41,8 @@ class LazyTreeContext(AbstractLazyContext): self._predefined_names = dict(context.predefined_names) def infer(self): - old, self._context.predefined_names = \ - self._context.predefined_names, self._predefined_names - try: + with monkeypatch(self._context, 'predefined_names', self._predefined_names): return self._context.eval_node(self.data) - finally: - self._context.predefined_names = old def get_merged_lazy_context(lazy_contexts):