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

aw_query/__init__.py

@@ -0,0 +1,3 @@
+from .query2 import query
+
+__all__ = ["query"]

aw_query/exceptions.py

@@ -0,0 +1,14 @@
+class QueryException(Exception):
+    pass
+
+
+class QueryFunctionException(QueryException):
+    pass
+
+
+class QueryParseException(QueryException):
+    pass
+
+
+class QueryInterpretException(QueryException):
+    pass

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)

aw_query/query2.py

@@ -0,0 +1,448 @@
+import logging
+from abc import abstractmethod
+from datetime import datetime
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+)
+
+from aw_datastore import Datastore
+
+from .exceptions import QueryInterpretException, QueryParseException
+from .functions import functions
+
+logger = logging.getLogger(__name__)
+
+
+class QToken:
+    @abstractmethod
+    def interpret(self, datastore: Datastore, namespace: dict):
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def parse(string: str, namespace: dict):
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def check(string: str) -> Tuple[Optional[str], str]:
+        raise NotImplementedError
+
+
+class QInteger(QToken):
+    def __init__(self, value) -> None:
+        self.value = value
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        return self.value
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> QToken:
+        return QInteger(int(string))
+
+    @staticmethod
+    def check(string: str):
+        token = ""
+        for char in string:
+            if char.isdigit():
+                token += char
+            else:
+                break
+        return token, string[len(token) :]
+
+
+class QVariable(QToken):
+    def __init__(self, name, value) -> None:
+        self.name = name
+        self.value = value
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        if self.name not in namespace:
+            raise QueryInterpretException(
+                f"Tried to reference variable '{self.name}' which is not defined"
+            )
+        namespace[self.name] = self.value
+        return self.value
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> QToken:
+        val = None
+        if string in namespace:
+            val = namespace[string]
+        return QVariable(string, val)
+
+    @staticmethod
+    def check(string: str):
+        token = ""
+        for i, char in enumerate(string):
+            if char.isalpha() or char == "_":
+                token += char
+            elif i != 0 and char.isdigit():
+                token += char
+            else:
+                break
+        return token, string[len(token) :]
+
+
+class QString(QToken):
+    def __init__(self, value):
+        self.value = value
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        return self.value
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> "QString":
+        quotes_type = string[0]
+        string = string.replace("\\" + quotes_type, quotes_type)
+        string = string[1:-1]
+        return QString(string)
+
+    @staticmethod
+    def check(string: str):
+        token = ""
+        quotes_type = string[0]
+        if quotes_type != '"' and quotes_type != "'":
+            return token, string
+        token += quotes_type
+        prev_char = None
+        for char in string[1:]:
+            token += char
+            if (
+                char == quotes_type and prev_char != "\\"
+            ):  # escape quote_type with backslash
+                break
+            prev_char = char
+        if token[-1] != quotes_type or len(token) < 2:
+            # Unclosed string?
+            raise QueryParseException("Failed to parse string")
+        return token, string[len(token) :]
+
+
+class QFunction(QToken):
+    def __init__(self, name, args):
+        self.name = name
+        self.args = args
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        if self.name not in functions:
+            raise QueryInterpretException(
+                f"Tried to call function '{self.name}' which doesn't exist"
+            )
+        call_args = [datastore, namespace]
+        for arg in self.args:
+            call_args.append(arg.interpret(datastore, namespace))
+        # logger.debug("Arguments for functioncall to {} is {}".format(self.name, call_args))
+        try:
+            result = functions[self.name](*call_args)  # type: ignore
+        except TypeError:
+            raise QueryInterpretException(
+                f"Tried to call function {self.name} with invalid amount of arguments"
+            ) from None
+        return result
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> QToken:
+        arg_start = 0
+        arg_end = len(string) - 1
+        # Find opening bracket
+        for char in string:
+            if char == "(":
+                break
+            arg_start = arg_start + 1
+        # Parse name
+        name = string[:arg_start]
+        # Parse arguments
+        args = []
+        args_str = string[arg_start + 1 : arg_end]
+        while args_str:
+            (arg_t, arg), args_str = _parse_token(args_str, namespace)
+            comma = args_str.find(",")
+            if comma != -1:
+                args_str = args_str[comma + 1 :]
+            args.append(arg_t.parse(arg, namespace))
+        return QFunction(name, args)
+
+    @staticmethod
+    def check(string: str):
+        i = 0
+        # Find opening bracket
+        found = False
+        for char in string:
+            if char.isalpha() or char == "_":
+                i = i + 1
+            elif i != 0 and char.isdigit():
+                i = i + 1
+            elif char == "(":
+                i = i + 1
+                found = True
+                break
+            else:
+                break
+        if not found:
+            return None, string
+        to_consume = 1
+        single_quote = False
+        double_quote = False
+        prev_char = None
+        for char in string[i:]:
+            i = i + 1
+            if char == "'" and prev_char != "\\" and not double_quote:
+                single_quote = not single_quote
+            elif char == '"' and prev_char != "\\" and not single_quote:
+                double_quote = not double_quote
+            elif single_quote or double_quote:
+                pass
+            elif i != 0 and char.isdigit():
+                pass
+            elif char == "(":
+                to_consume += 1
+            elif char == ")":
+                to_consume -= 1
+            if to_consume == 0:
+                break
+            prev_char = char
+        if to_consume != 0:
+            return None, string
+        return string[:i], string[i + 1 :]
+
+
+class QDict(QToken):
+    def __init__(self, value: dict) -> None:
+        self.value = value
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        expanded_dict = {}
+        for key, value in self.value.items():
+            expanded_dict[key] = value.interpret(datastore, namespace)
+        return expanded_dict
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> QToken:
+        entries_str = string[1:-1]
+        d: Dict[str, QToken] = {}
+        while len(entries_str) > 0:
+            entries_str = entries_str.strip()
+            if len(d) > 0 and entries_str[0] == ",":
+                entries_str = entries_str[1:]
+            # parse key
+            (key_t, key_str), entries_str = _parse_token(entries_str, namespace)
+            if key_t != QString:
+                raise QueryParseException("Key in dict is not a str")
+            key = QString.parse(key_str, {}).value
+            entries_str = entries_str.strip()
+            # Remove :
+            if entries_str[0] != ":":
+                raise QueryParseException("Key in dict is not followed by a :")
+            entries_str = entries_str[1:]
+            # parse val
+            (val_t, val_str), entries_str = _parse_token(entries_str, namespace)
+            if not val_t:
+                raise QueryParseException("Dict expected a value, got nothing")
+            val = val_t.parse(val_str, namespace)
+            # set
+            d[key] = val
+        return QDict(d)
+
+    @staticmethod
+    def check(string: str):
+        if string[0] != "{":
+            return None, string
+        # Find closing bracket
+        i = 1
+        to_consume = 1
+        single_quote = False
+        double_quote = False
+        prev_char = None
+        for char in string[i:]:
+            i += 1
+            if char == "'" and prev_char != "\\" and not double_quote:
+                single_quote = not single_quote
+            elif char == '"' and prev_char != "\\" and not single_quote:
+                double_quote = not double_quote
+            elif single_quote or double_quote:
+                pass
+            elif char == "}":
+                to_consume = to_consume - 1
+            elif char == "{":
+                to_consume = to_consume + 1
+            if to_consume == 0:
+                break
+            prev_char = char
+        return string[:i], string[i + 1 :]
+
+
+class QList(QToken):
+    def __init__(self, value: list) -> None:
+        self.value = value
+
+    def interpret(self, datastore: Datastore, namespace: dict):
+        expanded_list = []
+        for value in self.value:
+            expanded_list.append(value.interpret(datastore, namespace))
+        return expanded_list
+
+    @staticmethod
+    def parse(string: str, namespace: dict) -> QToken:
+        entries_str = string[1:-1]
+        ls: List[QToken] = []
+        while len(entries_str) > 0:
+            entries_str = entries_str.strip()
+            if len(ls) > 0 and entries_str[0] == ",":
+                entries_str = entries_str[1:]
+            # parse
+            (val_t, val_str), entries_str = _parse_token(entries_str, namespace)
+            if not val_t:
+                raise QueryParseException("List expected a value, got nothing")
+            val = val_t.parse(val_str, namespace)
+            # set
+            ls.append(val)
+        return QList(ls)
+
+    @staticmethod
+    def check(string: str):
+        if string[0] != "[":
+            return None, string
+        # Find closing bracket
+        i = 1
+        to_consume = 1
+        single_quote = False
+        double_quote = False
+        prev_char = None
+        for char in string[i:]:
+            i += 1
+            if char == "'" and prev_char != "\\" and not double_quote:
+                single_quote = not single_quote
+            elif char == '"' and prev_char != "\\" and not single_quote:
+                double_quote = not double_quote
+            elif double_quote or single_quote:
+                pass
+            elif char == "]":
+                to_consume = to_consume - 1
+            elif char == "[":
+                to_consume = to_consume + 1
+            if to_consume == 0:
+                break
+            prev_char = char
+        return string[:i], string[i + 1 :]
+
+
+qtypes: Sequence[Type[QToken]] = [QString, QInteger, QFunction, QDict, QList, QVariable]
+
+
+def _parse_token(string: str, namespace: dict) -> Tuple[Tuple[Any, str], str]:
+    # TODO: The whole parsing thing is shoddily written, needs a rewrite from ground-up
+    if not isinstance(string, str):
+        raise QueryParseException(
+            "Reached unreachable, cannot parse something that isn't a string"
+        )
+    if len(string) == 0:
+        return (None, ""), string
+    string = string.strip()
+    token = None
+    t = None  # Declare so we can return it
+    for t in qtypes:
+        token, string = t.check(string)
+        if token:
+            break
+    if not token:
+        raise QueryParseException(f"Syntax error: {string}")
+    return (t, token), string
+
+
+def create_namespace() -> dict:
+    namespace = {
+        "True": True,
+        "False": False,
+        "true": True,
+        "false": False,
+    }
+    return namespace
+
+
+def parse(line, namespace):
+    separator_i = line.find("=")
+    var_str = line[:separator_i]
+    val_str = line[separator_i + 1 :]
+    if not val_str:
+        # TODO: Proper message
+        raise QueryParseException("Nothing to assign")
+    (var_t, var), var_str = _parse_token(var_str, namespace)
+    var_str = var_str.strip()
+    if var_str:  # Didn't consume whole var string
+        raise QueryParseException("Invalid syntax for assignment variable")
+    if var_t is not QVariable:
+        raise QueryParseException("Cannot assign to a non-variable")
+    (val_t, val), var_str = _parse_token(val_str, namespace)
+    if var_str:  # Didn't consume whole val string
+        raise QueryParseException("Invalid syntax for value to assign")
+    # Parse token
+    var = var_t.parse(var, namespace)
+    val = val_t.parse(val, namespace)
+    return var, val
+
+
+def interpret(var, val, namespace, datastore):
+    namespace[var.name] = val.interpret(datastore, namespace)
+    # logger.debug("Set {} to {}".format(var.name, namespace[var.name]))
+
+
+def get_return(namespace):
+    if "RETURN" not in namespace:
+        raise QueryParseException(
+            "Query doesn't assign the RETURN variable, nothing to respond"
+        )
+    return namespace["RETURN"]
+
+
+def _split_query_statements(query: str) -> List[str]:
+    """Split query into statements on semicolons, ignoring semicolons inside string literals."""
+    statements = []
+    current = ""
+    in_single_quote = False
+    in_double_quote = False
+    prev_char = None
+
+    for char in query:
+        if char == "'" and prev_char != "\\" and not in_double_quote:
+            in_single_quote = not in_single_quote
+            current += char
+        elif char == '"' and prev_char != "\\" and not in_single_quote:
+            in_double_quote = not in_double_quote
+            current += char
+        elif char == ";" and not in_single_quote and not in_double_quote:
+            statements.append(current)
+            current = ""
+        else:
+            current += char
+        prev_char = char
+
+    if current:
+        statements.append(current)
+    return statements
+
+
+def query(
+    name: str, query: str, starttime: datetime, endtime: datetime, datastore: Datastore
+) -> Any:
+    namespace = create_namespace()
+    namespace["NAME"] = name
+    namespace["STARTTIME"] = starttime.isoformat()
+    namespace["ENDTIME"] = endtime.isoformat()
+
+    query_stmts = _split_query_statements(query)
+    for statement in query_stmts:
+        statement = statement.strip()
+        if statement:
+            logger.debug("Parsing: " + statement)
+            var, val = parse(statement, namespace)
+            interpret(var, val, namespace, datastore)
+
+    result = get_return(namespace)
+    return result