Add a generator cache for py__mro__

jedi/evaluate/cache.py

@@ -4,7 +4,10 @@
 - ``CachedMetaClass`` uses ``_memoize_default`` to do the same with classes.
 """
 
+from jedi import debug
+
 _NO_DEFAULT = object()
+_RECURSION_SENTINEL = object()
 
 
 def _memoize_default(default=_NO_DEFAULT, evaluator_is_first_arg=False, second_arg_is_evaluator=False):
@@ -28,8 +31,7 @@ def _memoize_default(default=_NO_DEFAULT, evaluator_is_first_arg=False, second_a
             try:
                 memo = cache[function]
             except KeyError:
-                memo = {}
-                cache[function] = memo
+                cache[function] = memo = {}
 
             key = (obj, args, frozenset(kwargs.items()))
             if key in memo:
@@ -75,3 +77,47 @@ class CachedMetaClass(type):
     @evaluator_as_method_param_cache()
     def __call__(self, *args, **kwargs):
         return super(CachedMetaClass, self).__call__(*args, **kwargs)
+
+
+def evaluator_method_generator_cache():
+    """
+    This is a special memoizer. It memoizes generators and also checks for
+    recursion errors and returns no further iterator elemends in that case.
+    """
+    def func(function):
+        def wrapper(obj, *args, **kwargs):
+            cache = obj.evaluator.memoize_cache
+            try:
+                memo = cache[function]
+            except KeyError:
+                cache[function] = memo = {}
+
+            key = (obj, args, frozenset(kwargs.items()))
+
+            if key in memo:
+                actual_generator, cached_lst = memo[key]
+            else:
+                actual_generator = function(obj, *args, **kwargs)
+                cached_lst = []
+                memo[key] = actual_generator, cached_lst
+
+            i = 0
+            while True:
+                try:
+                    next_element = cached_lst[i]
+                    if next_element is _RECURSION_SENTINEL:
+                        debug.warning('Found a generator recursion for %s' % obj)
+                        # This means we have hit a recursion.
+                        return
+                except IndexError:
+                    cached_lst.append(_RECURSION_SENTINEL)
+                    next_element = next(actual_generator, None)
+                    if next_element is None:
+                        cached_lst.pop()
+                        return
+                    cached_lst[-1] = next_element
+                yield next_element
+                i += 1
+        return wrapper
+
+    return func