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

aw_client/queries.py

@@ -0,0 +1,334 @@
+"""
+Common queries.
+
+Most of these are from: https://github.com/ActivityWatch/aw-webui/blob/master/src/queries.ts
+"""
+
+import dataclasses
+import json
+import re
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from typing import (
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+from typing_extensions import TypeGuard
+
+import aw_client
+
+from .classes import get_classes
+
+
+class EnhancedJSONEncoder(json.JSONEncoder):
+    """For encoding dataclasses into JSON"""
+
+    def default(self, o):
+        if dataclasses.is_dataclass(o):
+            return dataclasses.asdict(o)  # type: ignore
+        return super().default(o)
+
+
+"""
+Do these dataclasses look confusing?
+Read up on dataclass inheritance: https://stackoverflow.com/a/53085935/965332
+"""
+
+
+@dataclass
+class _QueryParamsDefaultsBase:
+    bid_browsers: List[str] = field(default_factory=list)
+    classes: List[Tuple[List[str], dict]] = field(default_factory=list)
+    filter_classes: List[List[str]] = field(default_factory=list)
+    filter_afk: bool = True
+    include_audible: bool = True
+
+
+@dataclass
+class QueryParams(_QueryParamsDefaultsBase):
+    pass
+
+
+@dataclass
+class _DesktopQueryParamsBase:
+    bid_window: str
+    bid_afk: str
+    always_active_pattern: Optional[str] = None
+
+
+@dataclass
+class DesktopQueryParams(QueryParams, _DesktopQueryParamsBase):
+    pass
+
+
+@dataclass
+class _AndroidQueryParamsBase:
+    bid_android: str
+
+
+@dataclass
+class AndroidQueryParams(QueryParams, _AndroidQueryParamsBase):
+    pass
+
+
+def isDesktopParams(params: QueryParams) -> TypeGuard[DesktopQueryParams]:
+    return isinstance(params, DesktopQueryParams)
+
+
+def isAndroidParams(params: QueryParams) -> TypeGuard[AndroidQueryParams]:
+    return isinstance(params, AndroidQueryParams)
+
+
+def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> str:
+    if not params.classes:
+        # if categories not explicitly set,
+        # get categories from server settings
+        params.classes = get_classes()
+
+    # Needs escaping for regex patterns like '\w' to work (JSON.stringify adds extra unnecessary escaping)
+    classes_str = json.dumps(params.classes, cls=EnhancedJSONEncoder)
+    classes_str = re.sub(r"\\\\", r"\\", classes_str)
+
+    cat_filter_str = json.dumps(params.filter_classes)
+
+    # For simplicity, we assume that bid_window and bid_android are exchangeable (note however it needs special treatment)
+    bid_window = (
+        params.bid_window
+        if isinstance(params, DesktopQueryParams)
+        else params.bid_android
+    )
+
+    return "\n".join(
+        [
+            # Fetch window/app events
+            f'events = flood(query_bucket(find_bucket("{bid_window}")));',
+            # On Android, merge events to avoid overload of events
+            (
+                'events = merge_events_by_keys(events, ["app"]);'
+                if isAndroidParams(params)
+                else ""
+            ),
+            # Fetch not-afk events
+            (
+                f"""
+            not_afk = flood(query_bucket(find_bucket("{params.bid_afk}")));
+            not_afk = filter_keyvals(not_afk, "status", ["not-afk"]);"""
+                + (
+                    """
+            not_treat_as_afk = filter_keyvals_regex(events, "app", "%s");
+            not_afk = period_union(not_afk, not_treat_as_afk);
+            not_treat_as_afk = filter_keyvals_regex(events, "title", "%s");
+            not_afk = period_union(not_afk, not_treat_as_afk);"""
+                    % (
+                        params.always_active_pattern.replace('"', '\\"'),
+                        params.always_active_pattern.replace('"', '\\"'),
+                    )
+                    if params.always_active_pattern
+                    else ""
+                )
+                if isDesktopParams(params)
+                else ""
+            ),
+            # Fetch browser events
+            (
+                (browserEvents(params) if isDesktopParams(params) else "")
+                + (  # Include focused and audible browser events as indications of not-afk
+                    """
+            audible_events = filter_keyvals(browser_events, "audible", [true]);
+            not_afk = period_union(not_afk, audible_events);
+            """
+                    if params.include_audible
+                    else ""
+                )
+                if params.bid_browsers
+                else ""
+            ),
+            # Filter out window events when the user was afk
+            (
+                "events = filter_period_intersect(events, not_afk);"
+                if isDesktopParams(params) and params.filter_afk
+                else ""
+            ),
+            # Categorize
+            f"events = categorize(events, {classes_str});" if params.classes else "",
+            # Filter out selected categories
+            (
+                f'events = filter_keyvals(events, "$category", {cat_filter_str});'
+                if params.filter_classes
+                else ""
+            ),
+        ]
+    )
+
+
+def pretty_query(query: str) -> str:
+    return "\n".join([line.strip() for line in query.split("\n") if line.strip()])
+
+
+def _browser_in_buckets(browser: str, browserbuckets: List[str]) -> Optional[str]:
+    for bucket in browserbuckets:
+        if browser in bucket:
+            return bucket
+    return None
+
+
+def browsersWithBuckets(browserbuckets: List[str]) -> List[Tuple[str, str]]:
+    """Returns a list of (browserName, bucketId) pairs for found browser buckets"""
+    browsername_to_bucketid: List[Tuple[str, Optional[str]]] = [
+        (browserName, _browser_in_buckets(browserName, browserbuckets))
+        for browserName in browser_appnames
+    ]
+
+    # Only return browsers for which a bucket could be found
+    return [t for t in browsername_to_bucketid if t[1]]  # type: ignore
+
+
+def browserEvents(params: DesktopQueryParams) -> str:
+    """Returns a list of active browser events (where the browser was the active window) from all browser buckets"""
+    code = "browser_events = [];"
+
+    for browserName, bucketId in browsersWithBuckets(params.bid_browsers):
+        browser_appnames_str = json.dumps(browser_appnames[browserName])
+        code += f"""
+          events_{browserName} = flood(query_bucket("{bucketId}"));
+          window_{browserName} = filter_keyvals(events, "app", {browser_appnames_str});
+          events_{browserName} = filter_period_intersect(events_{browserName}, window_{browserName});
+          events_{browserName} = split_url_events(events_{browserName});
+          browser_events = concat(browser_events, events_{browserName});
+          browser_events = sort_by_timestamp(browser_events);
+        """
+    return code
+
+
+browser_appnames = {
+    "chrome": [
+        # Chrome
+        "Google Chrome",
+        "Google-chrome",
+        "chrome.exe",
+        "google-chrome-stable",
+        # Chromium
+        "Chromium",
+        "Chromium-browser",
+        "Chromium-browser-chromium",
+        "chromium.exe",
+        # Pre-releases
+        "Google-chrome-beta",
+        "Google-chrome-unstable",
+        # Brave (should this be merged with the brave entry?)
+        "Brave-browser",
+    ],
+    "firefox": [
+        "Firefox",
+        "Firefox.exe",
+        "firefox",
+        "firefox.exe",
+        "Firefox Developer Edition",
+        "firefoxdeveloperedition",
+        "Firefox-esr",
+        "Firefox Beta",
+        "Nightly",
+        "org.mozilla.firefox",
+    ],
+    "opera": ["opera.exe", "Opera"],
+    "brave": ["brave.exe"],
+    "edge": [
+        "msedge.exe",  # Windows
+        "Microsoft Edge",  # macOS
+    ],
+    "vivaldi": ["Vivaldi-stable", "Vivaldi-snapshot", "vivaldi.exe"],
+}
+
+default_limit = 100
+
+
+def querystr_to_array(querystr: str) -> List[str]:
+    return [line + ";" for line in querystr.split(";") if line]
+
+
+def escape_doublequote(s: str) -> str:
+    return s.replace('"', '\\"')
+
+
+def fullDesktopQuery(
+    params: DesktopQueryParams,
+) -> str:
+    # Escape `"`
+    params.bid_window = escape_doublequote(params.bid_window)
+    params.bid_afk = escape_doublequote(params.bid_afk)
+    params.bid_browsers = [escape_doublequote(bucket) for bucket in params.bid_browsers]
+
+    # Build the base query
+    query = f"""
+    {canonicalEvents(params)}
+    title_events = sort_by_duration(merge_events_by_keys(events, ["app", "title"]));
+    app_events   = sort_by_duration(merge_events_by_keys(title_events, ["app"]));
+    cat_events   = sort_by_duration(merge_events_by_keys(events, ["$category"]));
+    app_events  = limit_events(app_events, {default_limit});
+    title_events  = limit_events(title_events, {default_limit});
+    duration = sum_durations(events);
+    """
+
+    # Add browser-related query parts if browser buckets exist
+    if params.bid_browsers:
+        query += f"""
+        browser_events = split_url_events(browser_events);
+        browser_urls = merge_events_by_keys(browser_events, ["url"]);
+        browser_urls = sort_by_duration(browser_urls);
+        browser_urls = limit_events(browser_urls, {default_limit});
+        browser_domains = merge_events_by_keys(browser_events, ["$domain"]);
+        browser_domains = sort_by_duration(browser_domains);
+        browser_domains = limit_events(browser_domains, {default_limit});
+        browser_duration = sum_durations(browser_events);
+        """
+    else:
+        query += """
+        browser_events = [];
+        browser_urls = [];
+        browser_domains = [];
+        browser_duration = 0;
+        """
+
+    # Add the return statement
+    query += """
+        RETURN = {
+            "events": events,
+            "window": {
+                "app_events": app_events,
+                "title_events": title_events,
+                "cat_events": cat_events,
+                "active_events": not_afk,
+                "duration": duration
+            },
+            "browser": {
+                "domains": browser_domains,
+                "urls": browser_urls,
+                "duration": browser_duration
+            }
+        };
+    """
+    return query
+
+
+def test_fullDesktopQuery():
+    params = DesktopQueryParams(
+        bid_window="aw-watcher-window_",
+        bid_afk="aw-watcher-afk_",
+    )
+    now = datetime.now(tz=timezone.utc)
+    start = now - timedelta(days=7)
+    end = now
+    timeperiods = [(start, end)]
+    query = fullDesktopQuery(params)
+
+    awc = aw_client.ActivityWatchClient("test")
+    res = awc.query(query, timeperiods)[0]
+    events = res["events"]
+    print(len(events))
+
+
+if __name__ == "__main__":
+    test_fullDesktopQuery()