1+ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+ # SPDX-License-Identifier: Apache-2.0
3+
4+ """
5+ Cache implementation for WRITEME to speed up repeated runs.
6+ """
7+
8+ import json
9+ import logging
10+ import os
11+ import pickle
12+ from pathlib import Path
13+ from typing import Any , Dict , Optional
14+
15+ logger = logging .getLogger (__name__ )
16+
17+ # Cache directory relative to the readmes directory
18+ CACHE_DIR = Path (__file__ ).parent / ".cache"
19+
20+
21+ def get_cache_enabled () -> bool :
22+ """Check if caching is enabled via environment variable."""
23+ return os .environ .get ("USE_METADATA_CACHE" , "0" ) == "1"
24+
25+
26+ def ensure_cache_dir () -> None :
27+ """Ensure the cache directory exists."""
28+ if not CACHE_DIR .exists ():
29+ CACHE_DIR .mkdir (exist_ok = True )
30+ logger .debug (f"Created cache directory: { CACHE_DIR } " )
31+
32+
33+ def get_cache_path (key : str ) -> Path :
34+ """Get the cache file path for a given key."""
35+ # Create a filename-safe version of the key
36+ safe_key = key .replace ("/" , "_" ).replace (":" , "_" )
37+ return CACHE_DIR / f"{ safe_key } .pickle"
38+
39+
40+ def save_to_cache (key : str , data : Any ) -> bool :
41+ """
42+ Save data to cache.
43+
44+ Args:
45+ key: Cache key
46+ data: Data to cache (must be pickle-able)
47+
48+ Returns:
49+ bool: True if successfully cached, False otherwise
50+ """
51+ if not get_cache_enabled ():
52+ return False
53+
54+ try :
55+ ensure_cache_dir ()
56+ cache_path = get_cache_path (key )
57+
58+ with open (cache_path , "wb" ) as f :
59+ pickle .dump (data , f )
60+
61+ logger .debug (f"Cached data for key: { key } " )
62+ return True
63+ except Exception as e :
64+ logger .warning (f"Failed to cache data for key { key } : { e } " )
65+ return False
66+
67+
68+ def load_from_cache (key : str ) -> Optional [Any ]:
69+ """
70+ Load data from cache.
71+
72+ Args:
73+ key: Cache key
74+
75+ Returns:
76+ The cached data or None if not found or caching disabled
77+ """
78+ if not get_cache_enabled ():
79+ return None
80+
81+ cache_path = get_cache_path (key )
82+
83+ if not cache_path .exists ():
84+ return None
85+
86+ try :
87+ with open (cache_path , "rb" ) as f :
88+ data = pickle .load (f )
89+
90+ logger .debug (f"Loaded data from cache for key: { key } " )
91+ return data
92+ except Exception as e :
93+ logger .warning (f"Failed to load cache for key { key } : { e } " )
94+ return None
95+
96+
97+ def clear_cache () -> None :
98+ """Clear all cached data."""
99+ if CACHE_DIR .exists ():
100+ for cache_file in CACHE_DIR .glob ("*.pickle" ):
101+ try :
102+ cache_file .unlink ()
103+ except Exception as e :
104+ logger .warning (f"Failed to delete cache file { cache_file } : { e } " )
105+
106+ logger .info ("Cache cleared" )
0 commit comments