diff --git a/aw_transform/classify.py b/aw_transform/classify.py index 898f416..7f3ba4f 100644 --- a/aw_transform/classify.py +++ b/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: diff --git a/tests/test_transforms.py b/tests/test_transforms.py index df627db..b7bf883 100644 --- a/tests/test_transforms.py +++ b/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)