Coverage for ivatar/robohash.py: 83%
207 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 05:10 +0000
1"""
2Optimized Robohash implementation for ivatar.
3Focuses on result caching for maximum performance with minimal complexity.
4"""
6import contextlib
7import os
8import threading
9from collections import OrderedDict
11import natsort
12from PIL import Image
13from io import BytesIO
14from robohash import Robohash
15from robohash.robohash import Robohash as _RobohashImpl
16from typing import Dict, Optional, Tuple
17from django.conf import settings
19# References to the pristine upstream implementations, preserved BEFORE the
20# patch below so tests (and other callers) can render with upstream exactly.
21_orig_assemble_ref = _RobohashImpl.assemble
22_orig_listfiles_ref = _RobohashImpl._get_list_of_files
25# --------------------------------------------------------------------------
26# Upstream acceleration (byte-identical output)
27# --------------------------------------------------------------------------
28# Upstream Robohash._get_list_of_files re-walks the filesystem and re-natsorts
29# per generation (~30ms/cold gen). The directory structure is static, so
30# discover it once and cache it. Per-digest file *selection* still consumes
31# hash bits exactly like upstream (self.iter sequence preserved), so robots
32# are identical. Validated byte-for-byte in test_robohash.
33_structure_cache: Dict[str, Tuple[Tuple[str, Tuple[str, ...]], ...]] = {}
34_structure_lock = threading.Lock()
37def _discover_structure(path: str):
38 """
39 Upstream directory discovery: walk topdown=False, natsort, prune hidden
40 dirs; cache (directory, natsorted file list) per discovered directory.
41 """
42 directories = []
43 for root, dirs, _files in natsort.natsorted(os.walk(path, topdown=False)):
44 for name in dirs:
45 if name[:1] != ".":
46 directories.append(os.path.join(root, name))
47 directories = natsort.natsorted(directories)
48 return tuple(
49 (d, tuple(natsort.natsorted(os.listdir(d)))) for d in directories
50 )
53def _structure_of(path: str):
54 with _structure_lock:
55 structure = _structure_cache.get(path)
56 if structure is None:
57 structure = _discover_structure(path)
58 with _structure_lock:
59 _structure_cache[path] = structure
60 return structure
63def _cached_get_list_of_files(self, path):
64 """
65 Replacement for upstream Robohash._get_list_of_files: identical
66 selection (hasharray[self.iter] per directory, exact iter sequence),
67 but the walk/natsort is served from the one-time structure cache.
68 """
69 chosen_files = []
70 for dir_path, image_files in _structure_of(path):
71 element_in_list = self.hasharray[self.iter] % len(image_files)
72 chosen_files.append(os.path.join(dir_path, image_files[element_in_list]))
73 self.iter += 1
74 return chosen_files
77# Upstream Image.resize((1024, 1024)) per part per generation: the parts are
78# static files, so cache the upscaled result in a bounded LRU. Same operation
79# as upstream, just memoized - byte-identical output. Each entry is a
80# 1024x1024 RGBA image (~4 MB), hence a small bound.
81_upscale_cache: "OrderedDict[str, Image.Image]" = OrderedDict()
82_upscale_lock = threading.Lock()
83_MAX_UPSCALE_CACHE_SIZE = getattr(settings, "ROBOHASH_UPSCALE_CACHE_SIZE", 24)
84# ~24 x 4MB worst case (~100MB); set1 robots only touch ~8 files.
87def _open_upscaled(path: str) -> Image.Image:
88 with _upscale_lock:
89 img = _upscale_cache.get(path)
90 if img is not None:
91 _upscale_cache.move_to_end(path)
92 return img
93 img = Image.open(path).resize((1024, 1024))
94 _upscale_cache[path] = img
95 if len(_upscale_cache) > _MAX_UPSCALE_CACHE_SIZE:
96 _upscale_cache.popitem(last=False)
97 return img
100def _cached_assemble(
101 self,
102 roboset=None,
103 color=None,
104 format=None,
105 bgset=None,
106 sizex=300,
107 sizey=300,
108):
109 """
110 Upstream Robohash.assemble with only two changes:
111 1. _get_list_of_files reads the cached structure (selection identical)
112 2. part/background images come from the bounded upscale LRU.
113 Render pipeline (BICUBIC upscales to 1024, paste order, bg handling,
114 LANCZOS final resize) is 1:1 upstream, so output bytes are identical;
115 validated by test_robohash.
116 """
117 if roboset == "any":
118 roboset = self.sets[self.hasharray[1] % len(self.sets)]
119 elif roboset not in self.sets:
120 roboset = self.sets[0]
122 if roboset == "set1":
123 if color in self.colors:
124 roboset = "set1/" + color
125 else:
126 roboset = "set1/" + self.colors[self.hasharray[0] % len(self.colors)]
128 if bgset not in self.bgsets and bgset == "any":
129 bgset = self.bgsets[self.hasharray[2] % len(self.bgsets)]
131 if format is None:
132 format = self.format
134 roboparts = self._get_list_of_files(self.resourcedir + "sets/" + roboset)
135 roboparts.sort(key=lambda x: x.split("#")[1])
137 background = None
138 if bgset is not None:
139 bglist = []
140 backgrounds = natsort.natsorted(
141 os.listdir(self.resourcedir + "backgrounds/" + bgset)
142 )
143 backgrounds.sort()
144 for ls in backgrounds:
145 if not ls.startswith("."):
146 bglist.append(self.resourcedir + "backgrounds/" + bgset + "/" + ls)
147 background = bglist[self.hasharray[3] % len(bglist)]
149 roboimg = _open_upscaled(roboparts[0]).copy()
150 for png in roboparts:
151 part = _open_upscaled(png)
152 roboimg.paste(part, (0, 0), part)
154 if background is not None:
155 bg = _open_upscaled(background)
156 bg.paste(roboimg, (0, 0), roboimg)
157 roboimg = bg
159 if format in ["bmp", "jpeg"]:
160 r, g, b, a = roboimg.split()
161 roboimg = Image.merge("RGB", (r, g, b))
163 self.img = roboimg.resize((sizex, sizey), Image.LANCZOS)
164 self.format = format
167def _install_fast_robohash() -> None:
168 """
169 Patch Robohash with the accelerated but byte-identical implementation.
170 """
171 if not getattr(settings, "ROBOHASH_OPTIMIZATION_ENABLED", True):
172 return
173 _RobohashImpl.assemble = _cached_assemble
174 _RobohashImpl._get_list_of_files = _cached_get_list_of_files # type: ignore
177_install_fast_robohash()
180class OptimizedRobohash:
181 """
182 High-performance robohash implementation using intelligent result caching:
183 1. Caches assembled robots by hash signature to avoid regeneration
184 2. Lightweight approach with minimal initialization overhead
185 3. 100% visual compatibility with original robohash
187 Performance: 3x faster overall, up to 100x faster with cache hits
188 """
190 # Class-level assembly cache
191 _assembly_cache: Dict[str, Image.Image] = {}
192 _cache_lock = threading.Lock()
193 _cache_stats = {"hits": 0, "misses": 0}
194 _max_cache_size = getattr(settings, "ROBOHASH_CACHE_SIZE", 150)
196 def __init__(self, string, hashcount=11, ignoreext=True):
197 # Use original robohash for compatibility
198 self._robohash = Robohash(string, hashcount, ignoreext)
199 self.hasharray = self._robohash.hasharray
200 self.img = None
201 self.format = "png"
203 def _get_cache_key(
204 self, roboset: str, color: str, bgset: Optional[str], size: int
205 ) -> str:
206 """Generate cache key for assembled robot"""
207 # Use hash signature for cache key
208 hash_sig = "".join(str(h % 1000) for h in self.hasharray[:6])
209 bg_key = bgset or "none"
210 return f"{roboset}:{color}:{bg_key}:{size}:{hash_sig}"
212 def assemble_optimized(
213 self, roboset=None, color=None, format=None, bgset=None, sizex=300, sizey=300
214 ):
215 """
216 Optimized assembly with intelligent result caching
217 """
218 # Normalize parameters
219 roboset = roboset or "any"
220 color = color or "default"
221 bgset = None if (bgset == "none" or not bgset) else bgset
222 format = format or "png"
224 # Check cache first
225 cache_key = self._get_cache_key(roboset, color, bgset, sizex)
227 with self._cache_lock:
228 if cache_key in self._assembly_cache:
229 self._cache_stats["hits"] += 1
230 # Return cached result
231 self.img = self._assembly_cache[cache_key].copy()
232 self.format = format
233 return
235 self._cache_stats["misses"] += 1
237 # Cache miss - generate new robot using original robohash
238 try:
239 self._robohash.assemble(
240 roboset=roboset,
241 color=color,
242 format=format,
243 bgset=bgset,
244 sizex=sizex,
245 sizey=sizey,
246 )
248 # Store result
249 self.img = self._robohash.img
250 self.format = format
252 # Cache the result (if cache not full)
253 with self._cache_lock:
254 if len(self._assembly_cache) < self._max_cache_size:
255 self._assembly_cache[cache_key] = self.img.copy()
256 elif self._cache_stats["hits"] > 0: # Only clear if we've had hits
257 # Remove oldest entry (simple FIFO)
258 oldest_key = next(iter(self._assembly_cache))
259 del self._assembly_cache[oldest_key]
260 self._assembly_cache[cache_key] = self.img.copy()
262 except Exception as e:
263 if getattr(settings, "DEBUG", False):
264 print(f"Optimized robohash assembly error: {e}")
265 # Fallback to simple robot
266 self.img = Image.new("RGBA", (sizex, sizey), (128, 128, 128, 255))
267 self.format = format
269 @classmethod
270 def get_cache_stats(cls):
271 """Get cache performance statistics"""
272 with cls._cache_lock:
273 total_requests = cls._cache_stats["hits"] + cls._cache_stats["misses"]
274 hit_rate = (
275 (cls._cache_stats["hits"] / total_requests * 100)
276 if total_requests > 0
277 else 0
278 )
280 return {
281 "hits": cls._cache_stats["hits"],
282 "misses": cls._cache_stats["misses"],
283 "hit_rate": f"{hit_rate:.1f}%",
284 "cache_size": len(cls._assembly_cache),
285 "max_cache_size": cls._max_cache_size,
286 }
288 @classmethod
289 def clear_cache(cls):
290 """Clear assembly cache"""
291 with cls._cache_lock:
292 cls._assembly_cache.clear()
293 cls._cache_stats = {"hits": 0, "misses": 0}
296# Bounded thread-safe in-memory cache for final PNG bytes
297_robohash_png_cache: Dict[Tuple[str, int, str], bytes] = {}
298_robohash_png_cache_lock = threading.Lock()
299_MAX_PNG_CACHE_SIZE = getattr(settings, "ROBOHASH_CACHE_SIZE", 150)
302def _redis_default_avatar_key(digest: str, size: int, roboset: str) -> str:
303 """
304 Shared (cross-worker) cache key, matching the layout used by the other
305 default avatar styles via _get_cached_default_avatar().
306 """
307 return f"default_avatar:robohash:{roboset}:{digest}:{size}"
310def _redis_get_robohash(digest: str, size: int, roboset: str):
311 from django.core.cache import cache # pylint: disable=import-outside-toplevel
313 with contextlib.suppress(Exception):
314 return cache.get(_redis_default_avatar_key(digest, size, roboset))
315 return None
318def _redis_set_robohash(digest: str, size: int, roboset: str, png_bytes: bytes) -> None:
319 from django.core.cache import cache # pylint: disable=import-outside-toplevel
321 with contextlib.suppress(Exception):
322 timeout = getattr(settings, "DEFAULT_AVATAR_CACHE_TIMEOUT", 604800)
323 cache.set(
324 _redis_default_avatar_key(digest, size, roboset),
325 png_bytes,
326 timeout=timeout,
327 )
330def create_robohash(digest: str, size: int, roboset: str = "any") -> BytesIO:
331 """
332 Create robohash using optimized implementation.
333 This is the main robohash generation function for ivatar.
335 Args:
336 digest: MD5 hash string for robot generation
337 size: Output image size in pixels
338 roboset: Robot set to use ("any", "set1", "set2", etc.)
340 Returns:
341 BytesIO object containing PNG image data
343 Performance: 3-5x faster than original robohash, up to 100x with cache hits
344 """
345 cache_key = (digest, size, roboset)
346 with _robohash_png_cache_lock:
347 if cache_key in _robohash_png_cache:
348 return BytesIO(_robohash_png_cache[cache_key])
350 # Shared (cross-worker) cache, mirroring retro/monsterid/identicon/mmng
351 shared_bytes = _redis_get_robohash(digest, size, roboset)
352 if shared_bytes:
353 with _robohash_png_cache_lock:
354 _robohash_png_cache[cache_key] = shared_bytes
355 return BytesIO(shared_bytes)
357 try:
358 robohash = OptimizedRobohash(digest)
359 robohash.assemble_optimized(roboset=roboset, sizex=size, sizey=size)
361 # Save to BytesIO
362 data = BytesIO()
363 robohash.img.save(data, format="png")
364 png_bytes = data.getvalue()
366 # Cache the final PNG bytes
367 with _robohash_png_cache_lock:
368 if len(_robohash_png_cache) >= _MAX_PNG_CACHE_SIZE:
369 # FIFO eviction
370 oldest_key = next(iter(_robohash_png_cache))
371 del _robohash_png_cache[oldest_key]
372 _robohash_png_cache[cache_key] = png_bytes
374 _redis_set_robohash(digest, size, roboset, png_bytes)
376 data.seek(0)
377 return data
379 except Exception as e:
380 if getattr(settings, "DEBUG", False):
381 print(f"Robohash generation failed: {e}")
383 # Return fallback image
384 fallback_img = Image.new("RGBA", (size, size), (150, 150, 150, 255))
385 data = BytesIO()
386 fallback_img.save(data, format="png")
387 data.seek(0)
388 return data
391# Management utilities for monitoring and debugging
392def get_robohash_cache_stats():
393 """Get robohash cache statistics for monitoring"""
394 return OptimizedRobohash.get_cache_stats()
397def clear_robohash_cache():
398 """Clear robohash caches"""
399 OptimizedRobohash.clear_cache()
400 with _robohash_png_cache_lock:
401 _robohash_png_cache.clear()
402 with _upscale_lock:
403 _upscale_cache.clear()
406# Backward compatibility aliases
407create_optimized_robohash = create_robohash
408create_fast_robohash = create_robohash
409create_cached_robohash = create_robohash