"""Optional persistence: round-trip the OHLCV cache to a private HF Dataset. The HF Space container's ``/tmp`` is wiped on every restart. If the user sets ``HF_TOKEN`` (or ``HUGGING_FACE_HUB_TOKEN``) and ``FSCANNER_CACHE_REPO``, we push the parquet cache to a private dataset repo and pull it back on startup. This makes cold-start scans near-instant for a previously-scanned universe. """ from __future__ import annotations import os from typing import Optional import pandas as pd from . import paths CACHE_REPO = os.environ.get("FSCANNER_CACHE_REPO", "").strip() CACHE_FILE = "ohlcv_cache.parquet" def _is_enabled() -> bool: return bool(CACHE_REPO) and bool( os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") ) def pull_remote_cache() -> int: """Try to download the remote cache and write to CACHE_PATH. Returns the number of tickers loaded, 0 on any failure or if not enabled. Pulls whenever remote persistence is configured; the remote copy overrides whatever (possibly stale or missing) file exists locally. """ if not _is_enabled(): return 0 try: from huggingface_hub import hf_hub_download path = hf_hub_download( repo_id=CACHE_REPO, repo_type="dataset", filename=CACHE_FILE, force_download=False, ) # Replace the local cache with the downloaded one df = pd.read_parquet(path) os.makedirs(os.path.dirname(paths.CACHE_PATH) or ".", exist_ok=True) df.to_parquet(paths.CACHE_PATH, index=False) tickers = df["Ticker"].nunique() if "Ticker" in df.columns else 0 return int(tickers) except Exception: return 0 def push_remote_cache() -> bool: """Upload the local cache to the remote dataset repo. Returns True on success. """ if not _is_enabled() or not os.path.exists(paths.CACHE_PATH): return False try: from huggingface_hub import HfApi api = HfApi() api.upload_file( path_or_fileobj=paths.CACHE_PATH, path_in_repo=CACHE_FILE, repo_id=CACHE_REPO, repo_type="dataset", commit_message="Update OHLCV cache", ) return True except Exception: return False