Coverage for ivatar/pagan_optimized.py: 78%

1""" 

2Optimized pagan avatar generator for ivatar 

3Provides 95x+ performance improvement through intelligent caching 

4""" 

5 

6import contextlib 

7import threading 

8from io import BytesIO 

9from typing import Dict, Optional, Tuple 

10from PIL import Image 

11from django.conf import settings 

12import pagan 

13 

14 

15def _fast_scale_pixels(color, layer): 

16 """ 

17 Drop-in replacement for pagan.generator.scale_pixels with identical 

18 output. 

19 

20 The upstream function iterates ALL IMAGE_SIZE x IMAGE_SIZE native pixels 

21 (128x128) for every layer and tests `in layer` (a list, O(n)) to find 

22 which virtual pixel a native pixel belongs to - millions of tuple 

23 comparisons per avatar. Only the VIRTUAL_RESOLUTION x VIRTUAL_RESOLUTION 

24 grid (16x16) can ever pass the bounds check, so iterate that directly 

25 and use a set for membership. Produces an identical pixelmap (same 

26 iteration order), validated byte-for-byte. 

27 """ 

28 pixelmap = [] 

29 layer_set = set(layer) 

30 ds_x = pagan.generator.dotsize[0] 

31 ds_y = pagan.generator.dotsize[1] 

32 max_px = int((pagan.generator.MAX_X - (ds_y - 1)) / ds_y) + 1 

33 max_py = int((pagan.generator.MAX_Y - (ds_x - 1)) / ds_x) + 1 

34 for pix_x in range(max_px): 

35 x1 = pix_x * ds_y 

36 x2 = x1 + (ds_y - 1) 

37 for pix_y in range(max_py): 

38 y1 = pix_y * ds_x 

39 y2 = y1 + (ds_x - 1) 

40 if (pix_x, pix_y) in layer_set: 

41 pixelmap.append([(y1, x1), (y2, x2), color]) 

42 return pixelmap 

43 

44 

45def _install_fast_scale_pixels() -> None: 

46 """ 

47 Patch pagan.generator.scale_pixels with the O(256) set-based version. 

48 Output is bit-identical (see test_pagan_optimized); only the iteration 

49 count drops (~115k -> ~2k per avatar). 

50 """ 

51 pagan.generator.scale_pixels = _fast_scale_pixels 

52 

53 

54_install_fast_scale_pixels() 

55 

56 

57class OptimizedPagan: 

58 """ 

59 Optimized pagan avatar generator that caches Avatar objects 

60 

61 Provides 95x+ performance improvement by caching expensive pagan.Avatar 

62 object creation while maintaining 100% visual compatibility 

63 """ 

64 

65 # Class-level cache shared across all instances 

66 _avatar_cache: Dict[str, pagan.Avatar] = {} 

67 _cache_lock = threading.Lock() 

68 _cache_stats = {"hits": 0, "misses": 0, "size": 0} 

69 

70 # Cache configuration 

71 _max_cache_size = getattr(settings, "PAGAN_CACHE_SIZE", 100) # Max cached avatars 

72 _cache_enabled = True # Always enabled - this is the default implementation 

73 

74 @classmethod 

75 def _get_cached_avatar(cls, digest: str) -> Optional[pagan.Avatar]: 

76 """Get cached pagan Avatar object or create and cache it""" 

77 

78 # Try to get from cache first 

79 with cls._cache_lock: 

80 if digest in cls._avatar_cache: 

81 cls._cache_stats["hits"] += 1 

82 return cls._avatar_cache[digest] 

83 

84 # Cache miss - create new Avatar object 

85 try: 

86 avatar = pagan.Avatar(digest) 

87 

88 with cls._cache_lock: 

89 # Cache management - remove oldest entries if cache is full 

90 if len(cls._avatar_cache) >= cls._max_cache_size: 

91 # Remove 20% of oldest entries to make room 

92 remove_count = max(1, cls._max_cache_size // 5) 

93 keys_to_remove = list(cls._avatar_cache.keys())[:remove_count] 

94 for key in keys_to_remove: 

95 del cls._avatar_cache[key] 

96 

97 # Cache the Avatar object 

98 cls._avatar_cache[digest] = avatar 

99 cls._cache_stats["misses"] += 1 

100 cls._cache_stats["size"] = len(cls._avatar_cache) 

101 

102 return avatar 

103 

104 except Exception as e: 

105 if getattr(settings, "DEBUG", False): 

106 print(f"Failed to create pagan avatar {digest}: {e}") 

107 return None 

108 

109 @classmethod 

110 def get_cache_stats(cls) -> Dict: 

111 """Get cache performance statistics""" 

112 with cls._cache_lock: 

113 total_requests = cls._cache_stats["hits"] + cls._cache_stats["misses"] 

114 hit_rate = ( 

115 (cls._cache_stats["hits"] / total_requests * 100) 

116 if total_requests > 0 

117 else 0 

118 ) 

119 

120 return { 

121 "size": cls._cache_stats["size"], 

122 "max_size": cls._max_cache_size, 

123 "hits": cls._cache_stats["hits"], 

124 "misses": cls._cache_stats["misses"], 

125 "hit_rate": f"{hit_rate:.1f}%", 

126 "total_requests": total_requests, 

127 } 

128 

129 @classmethod 

130 def clear_cache(cls): 

131 """Clear the pagan avatar cache (useful for testing or memory management)""" 

132 with cls._cache_lock: 

133 cls._avatar_cache.clear() 

134 cls._cache_stats = {"hits": 0, "misses": 0, "size": 0} 

135 

136 @classmethod 

137 def generate_optimized(cls, digest: str, size: int = 80) -> Optional[Image.Image]: 

138 """ 

139 Generate optimized pagan avatar 

140 

141 Args: 

142 digest (str): MD5 hash as hex string 

143 size (int): Output image size in pixels 

144 

145 Returns: 

146 PIL.Image: Resized pagan avatar image, or None on error 

147 """ 

148 try: 

149 # Get cached Avatar object (this is where the 95x speedup comes from) 

150 avatar = cls._get_cached_avatar(digest) 

151 if avatar is None: 

152 return None 

153 

154 # Resize the cached avatar's image (this is very fast ~0.2ms) 

155 # The original pagan avatar is 128x128 RGBA 

156 resized_img = avatar.img.resize((size, size), Image.LANCZOS) 

157 

158 return resized_img 

159 

160 except Exception as e: 

161 if getattr(settings, "DEBUG", False): 

162 print(f"Optimized pagan generation failed for {digest}: {e}") 

163 return None 

164 

165 

166# Bounded thread-safe in-memory cache for final PNG bytes 

167_pagan_png_cache: Dict[Tuple[str, int], bytes] = {} 

168_pagan_png_cache_lock = threading.Lock() 

169_MAX_PAGAN_PNG_CACHE_SIZE = getattr(settings, "PAGAN_CACHE_SIZE", 100) 

170 

171 

172def _redis_pagan_key(digest: str, size: int) -> str: 

173 """ 

174 Shared (cross-worker) cache key, matching the layout used by the other 

175 default avatar styles via _get_cached_default_avatar(). 

176 """ 

177 return f"default_avatar:pagan:{digest}:{size}" 

178 

179 

180def _redis_get_pagan(digest: str, size: int): 

181 from django.core.cache import cache # pylint: disable=import-outside-toplevel 

182 

183 with contextlib.suppress(Exception): 

184 return cache.get(_redis_pagan_key(digest, size)) 

185 return None 

186 

187 

188def _redis_set_pagan(digest: str, size: int, png_bytes: bytes) -> None: 

189 from django.core.cache import cache # pylint: disable=import-outside-toplevel 

190 

191 with contextlib.suppress(Exception): 

192 timeout = getattr(settings, "DEFAULT_AVATAR_CACHE_TIMEOUT", 604800) 

193 cache.set(_redis_pagan_key(digest, size), png_bytes, timeout=timeout) 

194 

195 

196def create_optimized_pagan(digest: str, size: int = 80) -> BytesIO: 

197 """ 

198 Create pagan avatar using optimized implementation 

199 Returns BytesIO object ready for HTTP response 

200 

201 Performance improvement: 95x+ faster than original pagan generation 

202 

203 Args: 

204 digest (str): MD5 hash as hex string 

205 size (int): Output image size in pixels 

206 

207 Returns: 

208 BytesIO: PNG image data ready for HTTP response 

209 """ 

210 cache_key = (digest, size) 

211 with _pagan_png_cache_lock: 

212 if cache_key in _pagan_png_cache: 

213 return BytesIO(_pagan_png_cache[cache_key]) 

214 

215 # Shared (cross-worker) cache, mirroring retro/monsterid/identicon/mmng 

216 shared_bytes = _redis_get_pagan(digest, size) 

217 if shared_bytes: 

218 with _pagan_png_cache_lock: 

219 _pagan_png_cache[cache_key] = shared_bytes 

220 return BytesIO(shared_bytes) 

221 

222 try: 

223 # Generate optimized pagan avatar 

224 img = OptimizedPagan.generate_optimized(digest, size) 

225 

226 if img is not None: 

227 # Save to BytesIO for HTTP response 

228 data = BytesIO() 

229 img.save(data, format="PNG") 

230 png_bytes = data.getvalue() 

231 

232 with _pagan_png_cache_lock: 

233 if len(_pagan_png_cache) >= _MAX_PAGAN_PNG_CACHE_SIZE: 

234 # Evict oldest 20% 

235 remove_count = max(1, _MAX_PAGAN_PNG_CACHE_SIZE // 5) 

236 keys_to_remove = list(_pagan_png_cache.keys())[:remove_count] 

237 for key in keys_to_remove: 

238 del _pagan_png_cache[key] 

239 _pagan_png_cache[cache_key] = png_bytes 

240 

241 _redis_set_pagan(digest, size, png_bytes) 

242 

243 data.seek(0) 

244 return data 

245 else: 

246 # Fallback to original implementation if optimization fails 

247 if getattr(settings, "DEBUG", False): 

248 print(f"Falling back to original pagan for {digest}") 

249 

250 paganobj = pagan.Avatar(digest) 

251 img = paganobj.img.resize((size, size), Image.LANCZOS) 

252 data = BytesIO() 

253 img.save(data, format="PNG") 

254 data.seek(0) 

255 return data 

256 

257 except Exception as e: 

258 if getattr(settings, "DEBUG", False): 

259 print(f"Pagan generation failed: {e}") 

260 

261 # Return simple fallback image on error 

262 fallback_img = Image.new("RGBA", (size, size), (100, 100, 150, 255)) 

263 data = BytesIO() 

264 fallback_img.save(data, format="PNG") 

265 data.seek(0) 

266 return data 

267 

268 

269# Management utilities 

270def get_pagan_cache_info(): 

271 """Get cache information for monitoring/debugging""" 

272 return OptimizedPagan.get_cache_stats() 

273 

274 

275def clear_pagan_cache(): 

276 """Clear the pagan avatar cache""" 

277 OptimizedPagan.clear_cache() 

278 

279 

280# Backward compatibility - maintain same interface as original 

281def create_pagan_avatar(digest: str, size: int = 80) -> BytesIO: 

282 """Backward compatibility alias for create_optimized_pagan""" 

283 return create_optimized_pagan(digest, size)