cache.py 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126
  1. """
  2. - the popular ``_memoize_default`` works like a typical memoize and returns the
  3. default otherwise.
  4. - ``CachedMetaClass`` uses ``_memoize_default`` to do the same with classes.
  5. """
  6. from functools import wraps
  7. from jedi import debug
  8. _NO_DEFAULT = object()
  9. _RECURSION_SENTINEL = object()
  10. def _memoize_default(default=_NO_DEFAULT, inference_state_is_first_arg=False,
  11. second_arg_is_inference_state=False):
  12. """ This is a typical memoization decorator, BUT there is one difference:
  13. To prevent recursion it sets defaults.
  14. Preventing recursion is in this case the much bigger use than speed. I
  15. don't think, that there is a big speed difference, but there are many cases
  16. where recursion could happen (think about a = b; b = a).
  17. """
  18. def func(function):
  19. def wrapper(obj, *args, **kwargs):
  20. # TODO These checks are kind of ugly and slow.
  21. if inference_state_is_first_arg:
  22. cache = obj.memoize_cache
  23. elif second_arg_is_inference_state:
  24. cache = args[0].memoize_cache # needed for meta classes
  25. else:
  26. cache = obj.inference_state.memoize_cache
  27. try:
  28. memo = cache[function]
  29. except KeyError:
  30. cache[function] = memo = {}
  31. key = (obj, args, frozenset(kwargs.items()))
  32. if key in memo:
  33. return memo[key]
  34. else:
  35. if default is not _NO_DEFAULT:
  36. memo[key] = default
  37. rv = function(obj, *args, **kwargs)
  38. memo[key] = rv
  39. return rv
  40. return wrapper
  41. return func
  42. def inference_state_function_cache(default=_NO_DEFAULT):
  43. def decorator(func):
  44. return _memoize_default(default=default, inference_state_is_first_arg=True)(func)
  45. return decorator
  46. def inference_state_method_cache(default=_NO_DEFAULT):
  47. def decorator(func):
  48. return _memoize_default(default=default)(func)
  49. return decorator
  50. def inference_state_as_method_param_cache():
  51. def decorator(call):
  52. return _memoize_default(second_arg_is_inference_state=True)(call)
  53. return decorator
  54. class CachedMetaClass(type):
  55. """
  56. This is basically almost the same than the decorator above, it just caches
  57. class initializations. Either you do it this way or with decorators, but
  58. with decorators you lose class access (isinstance, etc).
  59. """
  60. @inference_state_as_method_param_cache()
  61. def __call__(self, *args, **kwargs):
  62. return super().__call__(*args, **kwargs)
  63. def inference_state_method_generator_cache():
  64. """
  65. This is a special memoizer. It memoizes generators and also checks for
  66. recursion errors and returns no further iterator elemends in that case.
  67. """
  68. def func(function):
  69. @wraps(function)
  70. def wrapper(obj, *args, **kwargs):
  71. cache = obj.inference_state.memoize_cache
  72. try:
  73. memo = cache[function]
  74. except KeyError:
  75. cache[function] = memo = {}
  76. key = (obj, args, frozenset(kwargs.items()))
  77. if key in memo:
  78. actual_generator, cached_lst = memo[key]
  79. else:
  80. actual_generator = function(obj, *args, **kwargs)
  81. cached_lst = []
  82. memo[key] = actual_generator, cached_lst
  83. i = 0
  84. while True:
  85. try:
  86. next_element = cached_lst[i]
  87. if next_element is _RECURSION_SENTINEL:
  88. debug.warning('Found a generator recursion for %s' % obj)
  89. # This means we have hit a recursion.
  90. return
  91. except IndexError:
  92. cached_lst.append(_RECURSION_SENTINEL)
  93. next_element = next(actual_generator, None)
  94. if next_element is None:
  95. cached_lst.pop()
  96. return
  97. cached_lst[-1] = next_element
  98. yield next_element
  99. i += 1
  100. return wrapper
  101. return func