Source code for pycancensus.cache
"""
Caching functionality for pycancensus.
"""
import os
import json
import warnings
import pickle
import hashlib
from pathlib import Path
from typing import Any, Optional, List
import pandas as pd
import geopandas as gpd
from .settings import get_cache_path
# In-memory session cache for metadata (vector and region lists). Sits in
# front of the file cache so repeated calls within a session don't pay the
# disk read + unpickle cost on every access.
_session_cache: dict = {}
def session_cache_get(cache_key: str) -> Optional[Any]:
"""Get a value from the in-memory session cache."""
value = _session_cache.get(cache_key)
if isinstance(value, pd.DataFrame):
# Hand out copies so callers can't mutate the cached object
return value.copy()
return value
def session_cache_set(cache_key: str, value: Any) -> Any:
"""Store a value in the in-memory session cache."""
if isinstance(value, pd.DataFrame):
# Store a private copy so later caller mutations can't corrupt it
_session_cache[cache_key] = value.copy()
else:
_session_cache[cache_key] = value
return value
def _session_cache_remove(cache_keys: Optional[List[str]] = None) -> None:
"""Drop specific keys (or everything, if None) from the session cache."""
if cache_keys is None:
_session_cache.clear()
else:
for cache_key in cache_keys:
_session_cache.pop(cache_key, None)
def get_cached_data(cache_key: str) -> Optional[Any]:
"""
Retrieve data from cache if it exists.
Parameters
----------
cache_key : str
Unique identifier for the cached data.
Returns
-------
Any or None
Cached data if found, None otherwise.
"""
cache_path = Path(get_cache_path())
cache_file = cache_path / f"{cache_key}.pkl"
if cache_file.exists():
try:
with open(cache_file, "rb") as f:
return pickle.load(f)
except Exception:
# If cache file is corrupted, remove it
cache_file.unlink(missing_ok=True)
return None
def cache_data(cache_key: str, data: Any, metadata: Optional[dict] = None) -> None:
"""
Cache data to disk, optionally with a metadata sidecar.
Parameters
----------
cache_key : str
Unique identifier for the data.
data : Any
Data to cache.
metadata : dict, optional
Request metadata (dataset, level, vectors, data version, ...) stored
alongside the data in a .meta.json sidecar. Used for recalled-data
detection.
"""
cache_path = Path(get_cache_path())
cache_path.mkdir(parents=True, exist_ok=True)
cache_file = cache_path / f"{cache_key}.pkl"
try:
with open(cache_file, "wb") as f:
pickle.dump(data, f)
if metadata is not None:
meta_file = cache_path / f"{cache_key}.meta.json"
with open(meta_file, "w") as f:
json.dump(metadata, f)
except Exception as e:
warnings.warn(f"Failed to cache data: {e}")
def get_cache_metadata(cache_key: str) -> Optional[dict]:
"""Read the metadata sidecar for a cached entry, if present."""
meta_file = Path(get_cache_path()) / f"{cache_key}.meta.json"
if meta_file.exists():
try:
with open(meta_file) as f:
return json.load(f)
except Exception:
return None
return None
[docs]
def list_cache() -> pd.DataFrame:
"""
List all cached data files.
Returns
-------
pd.DataFrame
DataFrame with information about cached files including:
- cache_key: The cache key
- file_path: Full path to cached file
- size_mb: File size in MB
- created: Creation timestamp
- modified: Last modification timestamp
Examples
--------
>>> import pycancensus as pc
>>> cache_list = pc.list_cache()
>>> print(cache_list)
"""
cache_path = Path(get_cache_path())
if not cache_path.exists():
return pd.DataFrame(
columns=["cache_key", "file_path", "size_mb", "created", "modified"]
)
cache_files = []
for cache_file in cache_path.glob("*.pkl"):
try:
stat = cache_file.stat()
entry = {
"cache_key": cache_file.stem,
"file_path": str(cache_file),
"size_mb": round(stat.st_size / (1024 * 1024), 2),
"created": pd.Timestamp.fromtimestamp(stat.st_ctime),
"modified": pd.Timestamp.fromtimestamp(stat.st_mtime),
}
metadata = get_cache_metadata(cache_file.stem)
if metadata is not None:
for key in ("dataset", "level", "vectors", "version", "geo_version"):
entry[key] = metadata.get(key)
cache_files.append(entry)
except Exception:
continue
return pd.DataFrame(cache_files)
[docs]
def remove_from_cache(
cache_keys: Optional[List[str]] = None, all_cache: bool = False
) -> None:
"""
Remove items from cache.
Parameters
----------
cache_keys : list of str, optional
Specific cache keys to remove. If None and all_cache=False,
does nothing.
all_cache : bool, default False
If True, removes all cached data.
Examples
--------
>>> import pycancensus as pc
>>> # Remove specific cache entries
>>> pc.remove_from_cache(["regions_CA16", "vectors_CA16"])
>>>
>>> # Remove all cache (use with caution!)
>>> pc.remove_from_cache(all_cache=True)
"""
# Keep the in-memory session cache consistent with the file cache
if all_cache:
_session_cache_remove()
elif cache_keys:
_session_cache_remove(cache_keys)
cache_path = Path(get_cache_path())
if not cache_path.exists():
print("No cache directory found.")
return
removed_count = 0
if all_cache:
# Remove all .pkl files (and their metadata sidecars)
for cache_file in cache_path.glob("*.pkl"):
try:
cache_file.unlink()
(cache_path / f"{cache_file.stem}.meta.json").unlink(missing_ok=True)
removed_count += 1
except Exception as e:
print(f"Warning: Failed to remove {cache_file}: {e}")
print(f"Removed {removed_count} cached files.")
elif cache_keys:
# Remove specific cache keys
for cache_key in cache_keys:
cache_file = cache_path / f"{cache_key}.pkl"
if cache_file.exists():
try:
cache_file.unlink()
(cache_path / f"{cache_key}.meta.json").unlink(missing_ok=True)
removed_count += 1
print(f"Removed cache: {cache_key}")
except Exception as e:
print(f"Warning: Failed to remove {cache_key}: {e}")
else:
print(f"Cache key not found: {cache_key}")
if removed_count > 0:
print(f"Removed {removed_count} cached files.")
else:
print("No cache keys specified and all_cache=False. Nothing to remove.")
[docs]
def clear_cache() -> None:
"""
Clear all cached data.
This is an alias for remove_from_cache(all_cache=True).
Examples
--------
>>> import pycancensus as pc
>>> pc.clear_cache()
"""
remove_from_cache(all_cache=True)