# Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you under the Apache License, Version 2.0 (the# "License"); you may not use this file except in compliance# with the License. You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing,# software distributed under the License is distributed on an# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY# KIND, either express or implied. See the License for the# specific language governing permissions and limitations# under the License."""Memoize result of function via pickle, used for cache testcases."""# pylint: disable=broad-except,superfluous-parensimportosimportsysimportatexitfromdecoratorimportdecoratefrom.._ffi.baseimportstring_typestry:importcPickleaspickleexceptImportError:importpickle
[文档]classCache(object):"""A cache object for result cache. Parameters ---------- key: str The file key to the function save_at_exit: bool Whether save the cache to file when the program exits """cache_by_key={}def__init__(self,key,save_at_exit):cache_dir=".pkl_memoize_py{0}".format(sys.version_info[0])try:os.mkdir(cache_dir)exceptFileExistsError:passelse:self.cache={}self.path=os.path.join(cache_dir,key)ifos.path.exists(self.path):try:self.cache=pickle.load(open(self.path,"rb"))exceptException:self.cache={}else:self.cache={}self.dirty=Falseself.save_at_exit=save_at_exitdefsave(self):ifself.dirty:print("Save memoize result to %s"%self.path)withopen(self.path,"wb")asout_file:pickle.dump(self.cache,out_file,pickle.HIGHEST_PROTOCOL)
[文档]defmemoize(key,save_at_exit=False):"""Memoize the result of function and reuse multiple times. Parameters ---------- key: str The unique key to the file save_at_exit: bool Whether save the cache to file when the program exits Returns ------- fmemoize : function The decorator function to perform memoization. """def_register(f):"""Registration function"""allow_types=(string_types,int,float,tuple)fkey=key+"."+f.__name__+".pkl"iffkeynotinCache.cache_by_key:Cache.cache_by_key[fkey]=Cache(fkey,save_at_exit)cache=Cache.cache_by_key[fkey]cargs=tuple(x.cell_contentsforxinf.__closure__)iff.__closure__else()cargs=(len(cargs),)+cargsdef_memoized_f(func,*args,**kwargs):assertnotkwargs,"Only allow positional call"key=cargs+argsforarginkey:ifisinstance(arg,tuple):forxinarg:assertisinstance(x,allow_types)else:assertisinstance(arg,allow_types)ifkeyincache.cache:returncache.cache[key]res=func(*args)cache.cache[key]=rescache.dirty=Truereturnresreturndecorate(f,_memoized_f)return_register