xref: /aosp_15_r20/external/autotest/utils/frozen_chromite/utils/memoize.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# -*- coding: utf-8 -*-
2# Copyright 2018 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6"""Functions for automatic caching of expensive function calls."""
7
8from __future__ import print_function
9
10import functools
11import sys
12
13import six
14
15
16def MemoizedSingleCall(functor):
17  """Decorator for simple functor targets, caching the results
18
19  The functor must accept no arguments beyond either a class or self (depending
20  on if this is used in a classmethod/instancemethod context).  Results of the
21  wrapped method will be written to the class/instance namespace in a specially
22  named cached value.  All future invocations will just reuse that value.
23
24  Note that this cache is per-process, so sibling and parent processes won't
25  notice updates to the cache.
26  """
27  # TODO(build): Should we rebase to snakeoil.klass.cached* functionality?
28  # pylint: disable=protected-access
29  @functools.wraps(functor)
30  def wrapper(obj):
31    key = wrapper._cache_key
32    val = getattr(obj, key, None)
33    if val is None:
34      val = functor(obj)
35      setattr(obj, key, val)
36    return val
37
38  # Use name mangling to store the cached value in a (hopefully) unique place.
39  wrapper._cache_key = '_%s_cached' % (functor.__name__.lstrip('_'),)
40  return wrapper
41
42
43def Memoize(f):
44  """Decorator for memoizing a function.
45
46  Caches all calls to the function using a ._memo_cache dict mapping (args,
47  kwargs) to the results of the first function call with those args and kwargs.
48
49  If any of args or kwargs are not hashable, trying to store them in a dict will
50  cause a ValueError.
51
52  Note that this cache is per-process, so sibling and parent processes won't
53  notice updates to the cache.
54  """
55  # pylint: disable=protected-access
56  f._memo_cache = {}
57
58  @functools.wraps(f)
59  def wrapper(*args, **kwargs):
60    # Make sure that the key is hashable... as long as the contents of args and
61    # kwargs are hashable.
62    # TODO(phobbs) we could add an option to use the id(...) of an object if
63    # it's not hashable.  Then "MemoizedSingleCall" would be obsolete.
64    key = (tuple(args), tuple(sorted(kwargs.items())))
65    if key in f._memo_cache:
66      return f._memo_cache[key]
67
68    result = f(*args, **kwargs)
69    f._memo_cache[key] = result
70    return result
71
72  return wrapper
73
74
75def SafeRun(functors, combine_exceptions=False):
76  """Executes a list of functors, continuing on exceptions.
77
78  Args:
79    functors: An iterable of functors to call.
80    combine_exceptions: If set, and multiple exceptions are encountered,
81      SafeRun will raise a RuntimeError containing a list of all the exceptions.
82      If only one exception is encountered, then the default behavior of
83      re-raising the original exception with unmodified stack trace will be
84      kept.
85
86  Raises:
87    The first exception encountered, with corresponding backtrace, unless
88    |combine_exceptions| is specified and there is more than one exception
89    encountered, in which case a RuntimeError containing a list of all the
90    exceptions that were encountered is raised.
91  """
92  errors = []
93
94  for f in functors:
95    try:
96      f()
97    except Exception as e:
98      # Append the exception object and the traceback.
99      errors.append((e, sys.exc_info()[2]))
100
101  if errors:
102    if len(errors) == 1 or not combine_exceptions:
103      # To preserve the traceback.
104      inst, tb = errors[0]
105      six.reraise(inst, None, tb)
106    else:
107      raise RuntimeError([e[0] for e in errors])
108