wrf-python/src/wrf/cache.py

from __future__ import (absolute_import, division, print_function, 
                        unicode_literals)

from threading import local
from collections import OrderedDict

from .py3compat import py3range
from .config import get_cache_size

_local_storage = local()


def _shrink_cache():
    """Shrink the cache if applicable.
    
    This only applies if a user has modified the cache size, otherwise it 
    just returns.
    
    Returns:
    
        None
        
    """
    global _local_storage
    
    try:
        cache = _local_storage.cache
    except AttributeError:
        return
    
    diff = len(cache) - get_cache_size()
    
    if diff > 0:
        for _ in py3range(diff):
            cache.popitem(last=False)
            

def cache_item(key, product, value):
    """Store an item in the threadlocal cache.
    
    The cache should be viewed as two nested dictionaries.  The outer key is 
    usually the id for the sequence where the cached item was generated.  The 
    inner key is the product type.
    
    Storing a cached item behaves like:
    
        cache[key][product] = value
    
    The cache is thread local, so stored items are only available in 
    the thread that cached them.
    
    Args:
    
        key (:obj:`int`): The outer dictionary cache key, which is typically 
            the id of the sequence where the cached item was generated.
            
        product (:obj:`str`): The inner dictionary cache key, which is a 
            string for the product type.
        
        value (:obj:`object`): The object to store in the cache.
    
    Returns:
    
        None.
        
    See Also:
    
        :meth:`get_cached_item`
        
    """
    global _local_storage
    
    _shrink_cache()
    
    if key is None or get_cache_size() == 0:
        return
    
    try:
        cache = _local_storage.cache
    except AttributeError:
        _local_storage.cache = OrderedDict()
        cache = _local_storage.cache
    
    try:
        _ = cache[key]
    except KeyError:
        if len(cache) >= get_cache_size():
            cache.popitem(last=False) # Remove the oldest dataset
        
        cache[key] = OrderedDict()
        
    cache[key][product] = value
    
    
def get_cached_item(key, product):
    """Return an item from the threadlocal cache.
    
    The cache should be viewed as two nested dictionaries.  The outer key is 
    usually the id for the sequence where the cached item was generated.  The 
    inner key is the product type.
    
    Retrieving a cached item behaves like:
    
        value = cache[key][product]
    
    The cache is thread local, so stored items are only available in 
    the thread that cached them.
    
    Args:
    
        key (:obj:`int`): The outer dictionary cache key, which is typically 
            the id of the sequence where the cached item was generated.
            
        product (:obj:`str`): The inner dictionary cache key, which is a 
            string for the product type.
    
    Returns:
    
        :obj:`object`: The cached object.
        
    See Also:
    
        :meth:`cache_item`
        
    """
    global _local_storage
    
    _shrink_cache()
    
    if key is None or get_cache_size == 0:
        return None
    
    cache = getattr(_local_storage, "cache", None)
    
    if cache is None:
        return None
    
    if len(cache) == 0:
        return None
    
    prod_dict = cache.get(key, None)
    
    if prod_dict is None:
        return None
    
    result = prod_dict.get(product, None)
        
    return result


def _get_cache():
    """Return the threadlocal cache.
    
    This is primarily used for testing.
    
    Returns:
    
        :class:`threading.local`
    
    """
    global _local_storage
    
    _shrink_cache()
    return getattr(_local_storage, "cache", None)