Cache decorator for Python

cache.py

""" 
Basic cache module for Python 
See main part for usage example.
Author: Piotr Likus, 2024
License: BSD
"""

import os.path
import pickle 
import time
import hashlib
import functools 


class CacheBackend:
    def value_exists(self, key: str):
        return False
    
    def write(self, key: str, value):
        pass
    
    def read(self, key: str):
        return None 
        
    def clear(self):
        pass
        
    def invalidate_by_key(self, key: str):
        pass
    
class MemoryCacheBackend(CacheBackend):
    def __init__(self):
        self.data = dict()

    def value_exists(self, key: str):
        return key in self.data
    
    def write(self, key: str, value):
        self.data[key] = value
    
    def read(self, key: str):
        return self.data[key] 

    def clear(self):
        self.data.clear()
    
    def invalidate_by_key(self, key: str):
        if key in self.data:
            del self.data[key]
        
class FileCacheBackend(CacheBackend):
    def __init__(self, cache_dir):
        self.cache_dir = cache_dir 

    def value_exists(self, key: str):
        file_path = self.build_file_path(key)
        return os.path.isfile(file_path)

    def invalidate_by_key(self, key: str):
        file_path = self.build_file_path(key)
        if os.path.isfile(file_path):
            os.remove(file_path)
    
    def write(self, key: str, value):
        if not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)        
        
        file_path = self.build_file_path(key)    
        self.write_value_to_file(file_path, value)

    def read(self, key: str):
        file_path = self.build_file_path(key)    
        return self.read_value_from_file(file_path)

    def clear(self):
        file_ext = self.file_ext()
        file_list = os.listdir(self.cache_dir)
        
        for fname in file_list:
            if fname.endswith(file_ext):
                os.remove(os.path.join(self.cache_dir, fname))
            
    def read_value_from_file(self, file_path):

        with open(file_path, 'rb') as f:
            value_loaded = pickle.load(f) 
            
        return value_loaded
        
    def write_value_to_file(self, file_path, value):
        with open(file_path, 'wb') as f:  # open a text file        
            pickle.dump(value, f)
            
    def build_file_path(self, key: str):
        new_key = "".join([ c if c.isalnum() else "_" for c in key ])
        new_key = new_key.replace(" ", "_")
        file_name = new_key + self.file_ext()
        return os.path.join(self.cache_dir, file_name)
        
    def file_ext(self):
        return ".pkl"
            
def init_cache_registry():
    global cache_backend_dict
    cache_backend_dict = dict()

def setup_cache(a_cache_backend, backend_name: str = "global"):
    global cache_backend_dict
    if cache_backend_dict == None:
        cache_backend_dict = dict()
    cache_backend_dict[backend_name] = a_cache_backend
    
def get_cache_backend(backend_name: str):
    return cache_backend_dict.get(backend_name)    
    
def cache_result(backend_name = "global", salt = None):    

    def inner_decorator(func):

        @functools.wraps(func)
        def wrapper(*args, **kwargs):               
            backend_obj = get_cache_backend(backend_name)
            key = (*args, *kwargs.items())
            key_str = str(key)
            
            if salt is not None:
                key_str += salt
            
            func_str = str(func.__name__)
            key_str += func_str    

            # -- basic solution - for in-memory cache only
            # -- string hash is changing value between sessions
            # key_hash = str(hash(key_str))
            
            hash_obj = hashlib.sha256(key_str.encode())
            key_hash = hash_obj.hexdigest()
            
            key = key_hash

            if backend_obj.value_exists(key):
                print("Retrieving result from cache...")
                return backend_obj.read(key)

            result = func(*args, **kwargs)
            backend_obj.write(key, result)

            return result

        return wrapper
    
    return inner_decorator

def cache_clear(backend_name = "global"):
    backend_obj = get_cache_backend(backend_name)
    backend_obj.clear()
    

@cache_result()
def heavy_computation(x) -> int:
    """ slow computation returning int """
    print(f"Calculating square of {x}")
    time.sleep(2)  
    return x * x

@cache_result()
def heavy_computation_2nd(x) -> int:
    """ slow computation returning int """
    print(f"Calculating sum of {x}")
    time.sleep(2)  
    return x + x

@cache_result("memory")
def dict_computation(x) -> dict:
    """ slow computation returning dict """
    print(f"Calculating text version of {x}")
    time.sleep(2)  
    return {"value": str(x), "items": [x, x+1]}

def dict_computation_2nd(x) -> dict:
    """ slow computation returning dict """
    print(f"Calculating text version of {x} with prefix")
    time.sleep(2)  
    return {"value": "@" + str(x)}
    
if __name__ == "__main__":
    init_cache_registry()
    
    file_backend = FileCacheBackend("cache")        
    setup_cache(file_backend)
    
    memory_backend = MemoryCacheBackend()
    setup_cache(memory_backend, "memory")
    
    # cache_clear() -> if you need a clean run   
    
    print("Performing test...")
    
    result = heavy_computation(3)
    print(f"Result #1 (cached on 2nd run) = {result}")

    result = heavy_computation(3)
    print(f"Result #2 (cached always) = {result}")

    result = heavy_computation(4)
    print(f"Result #3 (cached on 2nd run) = {result}")

    result = heavy_computation_2nd(3)
    print(f"Result #4 (cached on 2nd run) = {result}")
    
    result = dict_computation(5)
    print(f"Result #5 = {result}")

    result = dict_computation(5)
    print(f"Result #6 (cached always) = {result}")

    result = cache_result(salt = "abc")(dict_computation_2nd)(5)
    print(f"Result #7 (cached on second run, uses runtime key salt) = {result}")

    result = cache_result(salt = "efg")(dict_computation_2nd)(5)
    print(f"Result #8 (cached on second run, uses runtime key salt) = {result}")