perf(aw-transform): cache categorize/tag results per event data to fix month-view timeout

aw_transform/classify.py

@@ -1,3 +1,4 @@
+import json
 from typing import Pattern, List, Iterable, Tuple, Dict, Optional, Any
 from functools import reduce
 import re
@@ -43,7 +44,18 @@ def match(self, e: Event) -> bool:
 def categorize(
     events: List[Event], classes: List[Tuple[Category, Rule]]
 ) -> List[Event]:
-    return [_categorize_one(e, classes) for e in events]
+    cache: Dict[str, Category] = {}
+    for e in events:
+        try:
+            key = json.dumps(e.data, sort_keys=True)
+        except TypeError:
+            key = str(id(e.data))
+        if key not in cache:
+            cache[key] = _pick_category(
+                [_cls for _cls, rule in classes if rule.match(e)]
+            )
+        e.data["$category"] = list(cache[key])
+    return events
 
 
 def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event:
@@ -54,7 +66,16 @@ def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event:
 
 
 def tag(events: List[Event], classes: List[Tuple[Tag, Rule]]) -> List[Event]:
-    return [_tag_one(e, classes) for e in events]
+    cache: Dict[str, List[Tag]] = {}
+    for e in events:
+        try:
+            key = json.dumps(e.data, sort_keys=True)
+        except TypeError:
+            key = str(id(e.data))
+        if key not in cache:
+            cache[key] = [_cls for _cls, rule in classes if rule.match(e)]
+        e.data["$tags"] = list(cache[key])
+    return events
 
 
 def _tag_one(e: Event, classes: List[Tuple[Tag, Rule]]) -> Event:

tests/test_transforms.py

@@ -418,6 +418,36 @@ def test_categorize():
     assert events[3].data["$category"] == ["Uncategorized"]
 
 
+def test_categorize_cache_correctness():
+    """Cache reuses category for identical data; distinct data gets its own category."""
+    now = datetime.now(timezone.utc)
+
+    classes = [
+        (["Browser"], Rule({"regex": "Firefox"})),
+        (["Editor"], Rule({"regex": "vim"})),
+    ]
+    firefox_data = {"app": "Firefox", "title": "Home"}
+    vim_data = {"app": "vim", "title": "classify.py"}
+
+    # 50 Firefox events, 1 vim event, 50 more Firefox events
+    events = (
+        [Event(timestamp=now, duration=0, data=dict(firefox_data)) for _ in range(50)]
+        + [Event(timestamp=now, duration=0, data=dict(vim_data))]
+        + [Event(timestamp=now, duration=0, data=dict(firefox_data)) for _ in range(50)]
+    )
+    result = categorize(events, classes)
+
+    for e in result[:50]:
+        assert e.data["$category"] == ["Browser"]
+    assert result[50].data["$category"] == ["Editor"]
+    for e in result[51:]:
+        assert e.data["$category"] == ["Browser"]
+
+    # Mutating one event's category must not affect others sharing the same data fingerprint
+    result[0].data["$category"].append("MUTATED")
+    assert result[1].data["$category"] == ["Browser"]
+
+
 def test_tags():
     now = datetime.now(timezone.utc)