1# Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Caching utilities.""" 16 17import inspect 18import weakref 19 20 21# TODO(mdan): Add a garbage collection hook for cleaning up modules. 22class _TransformedFnCache(object): 23 """Generic hierarchical cache for transformed functions. 24 25 The keys are soft references (i.e. they are discarded when the key is 26 destroyed) created from the source function by `_get_key`. The subkeys are 27 strong references and can be any value. Typically they identify different 28 kinds of transformation. 29 """ 30 31 __slots__ = ('_cache',) 32 33 def __init__(self): 34 self._cache = weakref.WeakKeyDictionary() 35 36 def _get_key(self, entity): 37 raise NotImplementedError('subclasses must override') 38 39 def has(self, entity, subkey): 40 key = self._get_key(entity) 41 parent = self._cache.get(key, None) 42 if parent is None: 43 return False 44 return subkey in parent 45 46 def __getitem__(self, entity): 47 key = self._get_key(entity) 48 parent = self._cache.get(key, None) 49 if parent is None: 50 # The bucket is initialized to support this usage: 51 # cache[key][subkey] = value 52 self._cache[key] = parent = {} 53 return parent 54 55 def __len__(self): 56 return len(self._cache) 57 58 59class CodeObjectCache(_TransformedFnCache): 60 """A function cache based on code objects. 61 62 Code objects are good proxies for the source code of a function. 63 64 This cache efficiently handles functions that share code objects, such as 65 functions defined in a loop, bound methods, etc. 66 67 The cache falls back to the function object, if it doesn't have a code object. 68 """ 69 70 def _get_key(self, entity): 71 if hasattr(entity, '__code__'): 72 return entity.__code__ 73 else: 74 return entity 75 76 77class UnboundInstanceCache(_TransformedFnCache): 78 """A function cache based on unbound function objects. 79 80 Using the function for the cache key allows efficient handling of object 81 methods. 82 83 Unlike the _CodeObjectCache, this discriminates between different functions 84 even if they have the same code. This is needed for decorators that may 85 masquerade as another function. 86 """ 87 88 def _get_key(self, entity): 89 if inspect.ismethod(entity): 90 return entity.__func__ 91 return entity 92 93 94