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 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import inspect 22import weakref 23 24 25# TODO(mdan): Add a garbage collection hook for cleaning up modules. 26class _TransformedFnCache(object): 27 """Generic hierarchical cache for transformed functions. 28 29 The keys are soft references (i.e. they are discarded when the key is 30 destroyed) created from the source function by `_get_key`. The subkeys are 31 strong references and can be any value. Typically they identify different 32 kinds of transformation. 33 """ 34 35 __slots__ = ('_cache',) 36 37 def __init__(self): 38 self._cache = weakref.WeakKeyDictionary() 39 40 def _get_key(self, entity): 41 raise NotImplementedError('subclasses must override') 42 43 def has(self, entity, subkey): 44 key = self._get_key(entity) 45 parent = self._cache.get(key, None) 46 if parent is None: 47 return False 48 return subkey in parent 49 50 def __getitem__(self, entity): 51 key = self._get_key(entity) 52 parent = self._cache.get(key, None) 53 if parent is None: 54 # The bucket is initialized to support this usage: 55 # cache[key][subkey] = value 56 self._cache[key] = parent = {} 57 return parent 58 59 def __len__(self): 60 return len(self._cache) 61 62 63class CodeObjectCache(_TransformedFnCache): 64 """A function cache based on code objects. 65 66 Code objects are good proxies for the source code of a function. 67 68 This cache efficiently handles functions that share code objects, such as 69 functions defined in a loop, bound methods, etc. 70 71 The cache falls back to the function object, if it doesn't have a code object. 72 """ 73 74 def _get_key(self, entity): 75 if hasattr(entity, '__code__'): 76 return entity.__code__ 77 else: 78 return entity 79 80 81class UnboundInstanceCache(_TransformedFnCache): 82 """A function cache based on unbound function objects. 83 84 Using the function for the cache key allows efficient handling of object 85 methods. 86 87 Unlike the _CodeObjectCache, this discriminates between different functions 88 even if they have the same code. This is needed for decorators that may 89 masquerade as another function. 90 """ 91 92 def _get_key(self, entity): 93 if inspect.ismethod(entity): 94 return entity.__func__ 95 return entity 96 97 98