Skip to content

Commit a76de08

Browse files
feat(classify): optional priority field for category rules (#153)
* feat(classify): optional priority field for category rules Match ActivityWatch/aw-server-rust#663: integer `priority`/`weight` on rules, default rank is depth * 10 so values can slot between nesting levels. Existing configs keep their current ordering. * fix(classify): pick category without i64 fallback sentinel Replace the signed-64-bit Uncategorized rank with Optional[int] so a valid unbounded Python priority below -(2**63) still beats Uncategorized. Git-Session-Id: bob-monitor-aw-core-153
1 parent 81e4367 commit a76de08

4 files changed

Lines changed: 209 additions & 13 deletions

File tree

aw_query/functions.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -333,12 +333,18 @@ def q2_nop():
333333
@q2_function(categorize)
334334
@q2_typecheck
335335
def q2_categorize(events: list, classes: list):
336-
classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
336+
try:
337+
classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
338+
except ValueError as exc:
339+
raise QueryFunctionException(str(exc)) from None
337340
return categorize(events, classes)
338341

339342

340343
@q2_function(tag)
341344
@q2_typecheck
342345
def q2_tag(events: list, classes: list):
343-
classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
346+
try:
347+
classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
348+
except ValueError as exc:
349+
raise QueryFunctionException(str(exc)) from None
344350
return tag(events, classes)

aw_transform/classify.py

Lines changed: 40 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import json
22
from typing import Pattern, List, Iterable, Tuple, Dict, Optional, Any
3-
from functools import reduce
43
import re
54

65
from aw_core import Event
@@ -10,14 +9,29 @@
109
Category = List[str]
1110

1211

12+
def _parse_optional_priority(rules: Dict[str, Any]) -> Optional[int]:
13+
if "priority" in rules:
14+
val = rules["priority"]
15+
elif "weight" in rules:
16+
val = rules["weight"]
17+
else:
18+
return None
19+
# bool is a subclass of int
20+
if isinstance(val, bool) or not isinstance(val, int):
21+
raise ValueError("priority/weight must be an integer")
22+
return val
23+
24+
1325
class Rule:
1426
regex: Optional[Pattern]
1527
select_keys: Optional[List[str]]
1628
ignore_case: bool
29+
priority: Optional[int]
1730

1831
def __init__(self, rules: Dict[str, Any]) -> None:
1932
self.select_keys = rules.get("select_keys", None)
2033
self.ignore_case = rules.get("ignore_case", False)
34+
self.priority = _parse_optional_priority(rules)
2135

2236
# NOTE: Also checks that the regex isn't an empty string (which would erroneously match everything)
2337
regex_str = rules.get("regex", None)
@@ -52,15 +66,15 @@ def categorize(
5266
key = str(id(e.data))
5367
if key not in cache:
5468
cache[key] = _pick_category(
55-
[_cls for _cls, rule in classes if rule.match(e)]
69+
[(_cls, rule) for _cls, rule in classes if rule.match(e)]
5670
)
5771
e.data["$category"] = list(cache[key])
5872
return events
5973

6074

6175
def _categorize_one(e: Event, classes: List[Tuple[Category, Rule]]) -> Event:
6276
e.data["$category"] = _pick_category(
63-
[_cls for _cls, rule in classes if rule.match(e)]
77+
[(_cls, rule) for _cls, rule in classes if rule.match(e)]
6478
)
6579
return e
6680

@@ -83,11 +97,26 @@ def _tag_one(e: Event, classes: List[Tuple[Tag, Rule]]) -> Event:
8397
return e
8498

8599

86-
def _pick_category(tags: Iterable[Category]) -> Category:
87-
return reduce(_pick_deepest_cat, tags, ["Uncategorized"])
88-
89-
90-
def _pick_deepest_cat(t1: Category, t2: Category) -> Category:
91-
# t1 will be the accumulator when used in reduce
92-
# Always bias against t1, since it could be "Uncategorized"
93-
return t2 if len(t2) >= len(t1) else t1
100+
def _effective_rank(category: Category, rule: Rule) -> int:
101+
# Integer-only. Default is depth * 10 so explicit priorities can slot
102+
# between nesting levels (depth 1 → 10, depth 2 → 20). Relative order of
103+
# unprioritized rules is unchanged.
104+
if rule.priority is not None:
105+
return rule.priority
106+
return len(category) * 10
107+
108+
109+
def _pick_category(matches: Iterable[Tuple[Category, Rule]]) -> Category:
110+
category: Category = ["Uncategorized"]
111+
rank: Optional[int] = None
112+
for cat, rule in matches:
113+
if not cat:
114+
continue
115+
item_rank = _effective_rank(cat, rule)
116+
# None means no match yet, so any non-empty category wins — including
117+
# an explicit priority below a signed 64-bit floor. Equal ranks keep
118+
# the later match (same contract as the old depth-only `>=`).
119+
if rank is None or item_rank >= rank:
120+
category = cat
121+
rank = item_rank
122+
return category

tests/test_query2.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -338,6 +338,19 @@ def test_query2_merge_subwatcher_fields_invalid_conflict():
338338
query(qname, example_query, starttime, endtime, ds)
339339

340340

341+
def test_query2_categorize_invalid_priority():
342+
ds = mock_ds
343+
qname = "asd"
344+
starttime = iso8601.parse_date("1970-01-01")
345+
endtime = iso8601.parse_date("1970-01-02")
346+
example_query = """
347+
events = [];
348+
RETURN = categorize(events, [[["test"], {"regex": "test", "priority": "high"}]]);
349+
"""
350+
with pytest.raises(QueryFunctionException, match="integer"):
351+
query(qname, example_query, starttime, endtime, ds)
352+
353+
341354
@pytest.mark.parametrize("datastore", param_datastore_objects())
342355
def test_query2_function_in_function(datastore):
343356
qname = "asd"
@@ -614,6 +627,39 @@ def test_query2_query_categorize(datastore):
614627
datastore.delete_bucket(bid)
615628

616629

630+
@pytest.mark.parametrize("datastore", param_datastore_objects())
631+
def test_query2_query_categorize_priority(datastore):
632+
bid = "test_bucket_priority"
633+
qname = "test"
634+
starttime = iso8601.parse_date("1970")
635+
endtime = starttime + timedelta(hours=1)
636+
637+
example_query = rf"""
638+
events = query_bucket("{bid}");
639+
events = categorize(events, [
640+
[["A"], {{"regex": "test", "priority": 25}}],
641+
[["B", "B1"], {{"regex": "test"}}]
642+
]);
643+
RETURN = events;
644+
"""
645+
try:
646+
bucket = datastore.create_bucket(
647+
bucket_id=bid, type="test", client="test", hostname="test", name="asd"
648+
)
649+
bucket.insert(
650+
Event(
651+
data={"label": "test"},
652+
timestamp=starttime,
653+
duration=timedelta(seconds=1),
654+
)
655+
)
656+
result = query(qname, example_query, starttime, endtime, datastore)
657+
assert len(result) == 1
658+
assert result[0].data["$category"] == ["A"]
659+
finally:
660+
datastore.delete_bucket(bid)
661+
662+
617663
@pytest.mark.parametrize("datastore", param_datastore_objects())
618664
def test_query2_flood_with_pulsetime(datastore):
619665
"""Test that flood() accepts an optional pulsetime argument (fixes WorkReport query bug)."""

tests/test_transforms.py

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -418,6 +418,121 @@ def test_categorize():
418418
assert events[3].data["$category"] == ["Uncategorized"]
419419

420420

421+
def _event(value: str = "just a test") -> Event:
422+
return Event(timestamp=datetime.now(timezone.utc), duration=0, data={"key": value})
423+
424+
425+
def test_categorize_depth_wins_without_priority():
426+
events = categorize(
427+
[_event()],
428+
[
429+
(["A"], Rule({"regex": "test"})),
430+
(["B", "B1"], Rule({"regex": "test"})),
431+
],
432+
)
433+
assert events[0].data["$category"] == ["B", "B1"]
434+
435+
436+
def test_categorize_explicit_priority_overrides_depth():
437+
# Default for B1 is depth 2 → 20; 25 beats it.
438+
events = categorize(
439+
[_event()],
440+
[
441+
(["A"], Rule({"regex": "test", "priority": 25})),
442+
(["B", "B1"], Rule({"regex": "test"})),
443+
],
444+
)
445+
assert events[0].data["$category"] == ["A"]
446+
447+
448+
def test_categorize_weight_alias():
449+
events = categorize(
450+
[_event()],
451+
[
452+
(["A"], Rule({"regex": "test", "weight": 25})),
453+
(["B", "B1"], Rule({"regex": "test"})),
454+
],
455+
)
456+
assert events[0].data["$category"] == ["A"]
457+
458+
459+
def test_categorize_inter_level_priority():
460+
between = categorize(
461+
[_event()],
462+
[
463+
(["A"], Rule({"regex": "test"})),
464+
(["A2"], Rule({"regex": "test", "priority": 15})),
465+
],
466+
)
467+
assert between[0].data["$category"] == ["A2"]
468+
469+
still_loses_to_deeper = categorize(
470+
[_event()],
471+
[
472+
(["A2"], Rule({"regex": "test", "priority": 15})),
473+
(["B", "B1"], Rule({"regex": "test"})),
474+
],
475+
)
476+
assert still_loses_to_deeper[0].data["$category"] == ["B", "B1"]
477+
478+
479+
def test_categorize_lower_priority_loses_to_default_depth():
480+
events = categorize(
481+
[_event()],
482+
[
483+
(["A"], Rule({"regex": "test"})),
484+
(["B", "B1"], Rule({"regex": "test", "priority": 0})),
485+
],
486+
)
487+
assert events[0].data["$category"] == ["A"]
488+
489+
490+
def test_categorize_equal_priority_keeps_later_match():
491+
events = categorize(
492+
[_event()],
493+
[
494+
(["First"], Rule({"regex": "test", "priority": 5})),
495+
(["Second"], Rule({"regex": "test", "priority": 5})),
496+
],
497+
)
498+
assert events[0].data["$category"] == ["Second"]
499+
500+
501+
def test_categorize_negative_priority_still_beats_uncategorized():
502+
events = categorize(
503+
[_event()],
504+
[(["Low"], Rule({"regex": "test", "priority": -100}))],
505+
)
506+
assert events[0].data["$category"] == ["Low"]
507+
508+
509+
def test_categorize_priority_below_i64_min_still_beats_uncategorized():
510+
# Python ints are unbounded; a signed-64-bit fallback sentinel would
511+
# incorrectly keep Uncategorized for values below -(2**63).
512+
events = categorize(
513+
[_event()],
514+
[(["Low"], Rule({"regex": "test", "priority": -(2**63) - 1}))],
515+
)
516+
assert events[0].data["$category"] == ["Low"]
517+
518+
519+
def test_categorize_empty_category_keeps_uncategorized():
520+
events = categorize(
521+
[_event()],
522+
[([], Rule({"regex": "test"}))],
523+
)
524+
assert events[0].data["$category"] == ["Uncategorized"]
525+
526+
527+
def test_rule_invalid_priority():
528+
with pytest.raises(ValueError, match="integer"):
529+
Rule({"regex": "test", "priority": 1.5})
530+
with pytest.raises(ValueError, match="integer"):
531+
Rule({"regex": "test", "priority": "high"})
532+
with pytest.raises(ValueError, match="integer"):
533+
Rule({"regex": "test", "priority": True})
534+
535+
421536
def test_categorize_cache_correctness():
422537
"""Cache reuses category for identical data; distinct data gets its own category."""
423538
now = datetime.now(timezone.utc)

0 commit comments

Comments
 (0)