snapshot of https://github.com/ActivityWatch/aw-core.git (master, 2026-09-16; blobs > 131072B omitted)

aw_query/functions.py

@@ -0,0 +1,350 @@
+from datetime import timedelta
+from functools import wraps
+from inspect import signature
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+)
+
+import iso8601
+from aw_core.models import Event
+from aw_datastore import Datastore
+from aw_transform import (
+    Rule,
+    categorize,
+    chunk_events_by_key,
+    concat,
+    filter_keyvals,
+    filter_keyvals_regex,
+    filter_period_intersect,
+    flood,
+    limit_events,
+    merge_events_by_keys,
+    merge_subwatcher_fields,
+    period_union,
+    simplify_string,
+    sort_by_duration,
+    sort_by_timestamp,
+    split_url_events,
+    sum_durations,
+    tag,
+    union_no_overlap,
+)
+
+from .exceptions import QueryFunctionException
+
+
+def _verify_bucket_exists(datastore, bucketname):
+    if bucketname in datastore.buckets():
+        return
+    else:
+        raise QueryFunctionException(f"There's no bucket named '{bucketname}'")
+
+
+def _verify_variable_is_type(variable, t):
+    if not isinstance(variable, t):
+        raise QueryFunctionException(
+            f"Variable '{variable}' passed to function call is of invalid type. Expected {t} but was {type(variable)}"
+        )
+
+
+# TODO: proper type checking (typecheck-decorator in pypi?)
+
+
+TNamespace = Dict[str, Any]
+TQueryFunction = Callable[..., Any]
+
+
+"""
+    Declarations
+"""
+functions: Dict[str, TQueryFunction] = {}
+
+
+def q2_function(transform_func=None):
+    """
+    Decorator used to register query functions.
+
+    Automatically adds mock arguments for Datastore and TNamespace
+    if not in function signature.
+    """
+
+    def h(f):
+        sig = signature(f)
+        # If function lacks docstring, use docstring from underlying function in aw_transform
+        if transform_func and transform_func.__doc__ and not f.__doc__:
+            f.__doc__ = f".. note:: Documentation automatically copied from underlying function `aw_transform.{transform_func.__name__}`\n\n{transform_func.__doc__}"
+
+        @wraps(f)
+        def g(datastore: Datastore, namespace: TNamespace, *args, **kwargs):
+            # Remove datastore and namespace argument for functions that don't need it
+            args = (datastore, namespace, *args)
+            if TNamespace not in (sig.parameters[p].annotation for p in sig.parameters):
+                args = (args[0], *args[2:])
+            if Datastore not in (sig.parameters[p].annotation for p in sig.parameters):
+                args = args[1:]
+            return f(*args, **kwargs)
+
+        fname = f.__name__
+        if fname[:3] == "q2_":
+            fname = fname[3:]
+        functions[fname] = g
+        return g
+
+    return h
+
+
+def q2_typecheck(f):
+    """Decorator that typechecks using `_verify_variable_is_type`"""
+    sig = signature(f)
+
+    @wraps(f)
+    def g(*args, **kwargs):
+        # FIXME: If the first argument passed to a query2 function is a straight [] then the second argument disappears from the argument list for unknown reasons, which breaks things
+        for i, p in enumerate(sig.parameters):
+            param = sig.parameters[p]
+
+            # print(f"Checking that param ({param}) was {param.annotation}, value: {args[i]}")
+            # FIXME: Won't check keyword arguments
+            if (
+                param.annotation in [list, str, int, float]
+                and param.default == param.empty
+            ):
+                _verify_variable_is_type(args[i], param.annotation)
+
+        return f(*args, **kwargs)
+
+    return g
+
+
+"""
+    Getting buckets
+"""
+
+
+@q2_function()
+@q2_typecheck
+def q2_find_bucket(
+    datastore: Datastore, filter_str: str, hostname: Optional[str] = None
+):
+    """Find bucket by using a filter_str (to avoid hardcoding bucket names)"""
+    for bucket in datastore.buckets():
+        if filter_str in bucket:
+            bucket_metadata = datastore[bucket].metadata()
+            if hostname:
+                if bucket_metadata["hostname"] == hostname:
+                    return bucket
+            else:
+                return bucket
+    raise QueryFunctionException(
+        f"Unable to find bucket matching '{filter_str}' (hostname filter set to '{hostname}')"
+    )
+
+
+"""
+    Data gathering functions
+"""
+
+
+@q2_function()
+@q2_typecheck
+def q2_query_bucket(
+    datastore: Datastore, namespace: TNamespace, bucketname: str
+) -> List[Event]:
+    _verify_bucket_exists(datastore, bucketname)
+    try:
+        starttime = iso8601.parse_date(namespace["STARTTIME"])
+        endtime = iso8601.parse_date(namespace["ENDTIME"])
+    except iso8601.ParseError:
+        raise QueryFunctionException(
+            "Unable to parse starttime/endtime for query_bucket"
+        ) from None
+    return datastore[bucketname].get(starttime=starttime, endtime=endtime)
+
+
+@q2_function()
+@q2_typecheck
+def q2_query_bucket_eventcount(
+    datastore: Datastore, namespace: TNamespace, bucketname: str
+) -> int:
+    _verify_bucket_exists(datastore, bucketname)
+    starttime = iso8601.parse_date(namespace["STARTTIME"])
+    endtime = iso8601.parse_date(namespace["ENDTIME"])
+    return datastore[bucketname].get_eventcount(starttime=starttime, endtime=endtime)
+
+
+"""
+    Filtering functions
+"""
+
+
+@q2_function(filter_keyvals)
+@q2_typecheck
+def q2_filter_keyvals(events: list, key: str, vals: list) -> List[Event]:
+    return filter_keyvals(events, key, vals, False)
+
+
+@q2_function(filter_keyvals)
+@q2_typecheck
+def q2_exclude_keyvals(events: list, key: str, vals: list) -> List[Event]:
+    return filter_keyvals(events, key, vals, True)
+
+
+@q2_function(filter_keyvals_regex)
+@q2_typecheck
+def q2_filter_keyvals_regex(events: list, key: str, regex: str) -> List[Event]:
+    return filter_keyvals_regex(events, key, regex)
+
+
+@q2_function(filter_period_intersect)
+@q2_typecheck
+def q2_filter_period_intersect(events: list, filterevents: list) -> List[Event]:
+    return filter_period_intersect(events, filterevents)
+
+
+@q2_function(period_union)
+@q2_typecheck
+def q2_period_union(events1: list, events2: list) -> List[Event]:
+    return period_union(events1, events2)
+
+
+@q2_function(limit_events)
+@q2_typecheck
+def q2_limit_events(events: list, count: int) -> List[Event]:
+    return limit_events(events, count)
+
+
+"""
+    Merge functions
+"""
+
+
+@q2_function(merge_events_by_keys)
+@q2_typecheck
+def q2_merge_events_by_keys(events: list, keys: list) -> List[Event]:
+    return merge_events_by_keys(events, keys)
+
+
+@q2_function(merge_subwatcher_fields)
+@q2_typecheck
+def q2_merge_subwatcher_fields(
+    base_events: list, subwatcher_events: list, keys: list, conflict: str = "base_wins"
+) -> List[Event]:
+    try:
+        return merge_subwatcher_fields(base_events, subwatcher_events, keys, conflict)
+    except ValueError as exc:
+        raise QueryFunctionException(str(exc)) from None
+
+
+@q2_function(chunk_events_by_key)
+@q2_typecheck
+def q2_chunk_events_by_key(events: list, key: str) -> List[Event]:
+    return chunk_events_by_key(events, key)
+
+
+"""
+    Sort functions
+"""
+
+
+@q2_function(sort_by_timestamp)
+@q2_typecheck
+def q2_sort_by_timestamp(events: list) -> List[Event]:
+    return sort_by_timestamp(events)
+
+
+@q2_function(sort_by_duration)
+@q2_typecheck
+def q2_sort_by_duration(events: list) -> List[Event]:
+    return sort_by_duration(events)
+
+
+"""
+    Summarizing functions
+"""
+
+
+@q2_function(sum_durations)
+@q2_typecheck
+def q2_sum_durations(events: list) -> timedelta:
+    return sum_durations(events)
+
+
+@q2_function(concat)
+@q2_typecheck
+def q2_concat(events1: list, events2: list) -> List[Event]:
+    return concat(events1, events2)
+
+
+@q2_function(union_no_overlap)
+@q2_typecheck
+def q2_union_no_overlap(events1: list, events2: list) -> List[Event]:
+    return union_no_overlap(events1, events2)
+
+
+"""
+    Flood functions
+"""
+
+
+@q2_function(flood)
+@q2_typecheck
+def q2_flood(events: list, pulsetime: float = 5) -> List[Event]:
+    return flood(events, pulsetime)
+
+
+"""
+    Watcher specific functions
+"""
+
+
+@q2_function(split_url_events)
+@q2_typecheck
+def q2_split_url_events(events: list) -> List[Event]:
+    return split_url_events(events)
+
+
+@q2_function(simplify_string)
+@q2_typecheck
+def q2_simplify_window_titles(events: list, key: str) -> List[Event]:
+    return simplify_string(events, key=key)
+
+
+"""
+    Test functions
+"""
+
+
+@q2_function()
+@q2_typecheck
+def q2_nop():
+    """No operation function for unittesting"""
+    return 1
+
+
+"""
+    Classify
+"""
+
+
+@q2_function(categorize)
+@q2_typecheck
+def q2_categorize(events: list, classes: list):
+    try:
+        classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    except ValueError as exc:
+        raise QueryFunctionException(str(exc)) from None
+    return categorize(events, classes)
+
+
+@q2_function(tag)
+@q2_typecheck
+def q2_tag(events: list, classes: list):
+    try:
+        classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
+    except ValueError as exc:
+        raise QueryFunctionException(str(exc)) from None
+    return tag(events, classes)