fix: fix fetching classes

aw_client/classes.py

@@ -4,16 +4,13 @@
 Taken from default classes in aw-webui
 """
 import logging
-import random
 from typing import (
     Any,
     Dict,
     List,
     Tuple,
 )
 
-import aw_client
-
 logger = logging.getLogger(__name__)
 
 CategoryId = List[str]
@@ -64,14 +61,11 @@
 ]
 
 
-def get_classes() -> List[Tuple[List[str], dict]]:
+def get_classes(awc) -> List[Tuple[List[str], dict]]:
     """
     Get classes from server-side settings.
     Might throw a 404 if not set yet, in which case we use the default classes as a fallback.
     """
-    # NOTE: Always tries to fetch from prod server,
-    #       which is potentially wrong if testing server is being used.
-    awc = aw_client.ActivityWatchClient(f"get-setting-{random.randint(0, 10000)}")
     try:
         classes = awc.get_setting("classes")
     except Exception:

aw_client/cli.py

@@ -1,20 +1,18 @@
 #!/usr/bin/env python3
-import json
 import argparse
+import json
 import logging
 import textwrap
-from typing import Optional, List
-from datetime import timedelta, datetime, timezone
+from datetime import datetime, timedelta, timezone
+from typing import List, Optional
 
 import click
-from tabulate import tabulate
-
 from aw_core import Event
+from tabulate import tabulate
 
 import aw_client
-from . import queries
-from .classes import default_classes
 
+from . import queries
 
 now = datetime.now(timezone.utc)
 td1day = timedelta(days=1)
@@ -161,6 +159,7 @@ def report(
     logger.info(f"Querying between {start} and {stop}")
     bid_window = f"aw-watcher-window_{hostname}"
     bid_afk = f"aw-watcher-afk_{hostname}"
+    awc = obj.client
 
     if not start.tzinfo:
         start = start.astimezone()
@@ -169,23 +168,18 @@ def report(
 
     bid_browsers: List[str] = []
 
-    # TODO: Allow loading from toml
-    logger.info("Using default classes")
-    classes = default_classes
-
     params = queries.DesktopQueryParams(
         bid_browsers=bid_browsers,
-        classes=classes,
         filter_classes=[],
         filter_afk=True,
         include_audible=True,
         bid_window=bid_window,
         bid_afk=bid_afk,
     )
-    query = queries.fullDesktopQuery(params)
+    query = queries.fullDesktopQuery(awc, params)
     logger.debug("Query: \n" + queries.pretty_query(query))
 
-    result = obj.client.query(query, [(start, stop)], cache=cache, name=name)
+    result = awc.query(query, [(start, stop)], cache=cache, name=name)
 
     # TODO: Print titles, apps, categories, with most time
     for period in result:
@@ -247,20 +241,19 @@ def canonical(
     logger.info(f"Querying between {start} and {stop}")
     bid_window = f"aw-watcher-window_{hostname}"
     bid_afk = f"aw-watcher-afk_{hostname}"
+    awc = obj.client
 
     if not start.tzinfo:
         start = start.astimezone()
     if not stop.tzinfo:
         stop = stop.astimezone()
 
-    classes = default_classes
-
     query = queries.canonicalEvents(
+        awc,
         queries.DesktopQueryParams(
             bid_window=bid_window,
             bid_afk=bid_afk,
-            classes=classes,
-        )
+        ),
     )
     query = f"""{query}\n RETURN = events;"""
     logger.debug("Query: \n" + queries.pretty_query(query))

aw_client/queries.py

@@ -80,11 +80,14 @@ def isAndroidParams(params: QueryParams) -> TypeGuard[AndroidQueryParams]:
     return isinstance(params, AndroidQueryParams)
 
 
-def canonicalEvents(params: Union[DesktopQueryParams, AndroidQueryParams]) -> str:
+def canonicalEvents(
+    awc: aw_client.ActivityWatchClient,
+    params: Union[DesktopQueryParams, AndroidQueryParams],
+) -> str:
     if not params.classes:
         # if categories not explicitly set,
         # get categories from server settings
-        params.classes = get_classes()
+        params.classes = get_classes(awc)
 
     # Needs escaping for regex patterns like '\w' to work (JSON.stringify adds extra unnecessary escaping)
     classes_str = json.dumps(params.classes, cls=EnhancedJSONEncoder)
@@ -233,6 +236,7 @@ def escape_doublequote(s: str) -> str:
 
 
 def fullDesktopQuery(
+    awc: aw_client.ActivityWatchClient,
     params: DesktopQueryParams,
 ) -> str:
     # Escape `"`
@@ -242,7 +246,7 @@ def fullDesktopQuery(
 
     return (
         f"""
-      {canonicalEvents(params)}
+      {canonicalEvents(awc, 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"]));
@@ -281,6 +285,7 @@ def fullDesktopQuery(
 
 
 def test_fullDesktopQuery():
+    awc = aw_client.ActivityWatchClient("test")
     params = DesktopQueryParams(
         bid_window="aw-watcher-window_",
         bid_afk="aw-watcher-afk_",
@@ -289,9 +294,8 @@ def test_fullDesktopQuery():
     start = now - timedelta(days=7)
     end = now
     timeperiods = [(start, end)]
-    query = fullDesktopQuery(params)
+    query = fullDesktopQuery(awc, params)
 
-    awc = aw_client.ActivityWatchClient("test")
     res = awc.query(query, timeperiods)[0]
     events = res["events"]
     print(len(events))

examples/load_dataframe.py

@@ -12,14 +12,15 @@
 from aw_client.queries import DesktopQueryParams, canonicalEvents
 
 
-def build_query() -> str:
+def build_query(awc) -> str:
     hostname = "fakedata" if os.getenv("CI") else socket.gethostname()
     canonicalQuery = canonicalEvents(
+        awc,
         DesktopQueryParams(
             bid_window=f"aw-watcher-window_{hostname}",
             bid_afk=f"aw-watcher-afk_{hostname}",
             classes=default_classes,
-        )
+        ),
     )
     return f"""
     {canonicalQuery}
@@ -31,10 +32,10 @@ def main() -> None:
     now = datetime.now(tz=timezone.utc)
     td30d = timedelta(days=30)
 
-    aw = ActivityWatchClient()
+    awc = ActivityWatchClient()
     print("Querying...")
-    query = build_query()
-    data = aw.query(query, [(now - td30d, now)])
+    query = build_query(awc)
+    data = awc.query(query, [(now - td30d, now)])
 
     events = [
         {

examples/suggest_categories.py

@@ -3,16 +3,19 @@
 
 This might make more sense as a notebook.
 """
-
 from collections import Counter
 from datetime import datetime, timedelta, timezone
-from tabulate import tabulate
-from typing import Dict, List, Tuple, Any
+from typing import (
+    Any,
+    Dict,
+    List,
+    Tuple,
+)
 
-from aw_core import Event
 import aw_client
 from aw_client import queries
-
+from aw_core import Event
+from tabulate import tabulate
 
 # set up client
 awc = aw_client.ActivityWatchClient("test")
@@ -38,11 +41,12 @@ def get_events(categories=List[Tuple[Tuple[str], Dict[str, Any]]]):
     timeperiods = [(start, now)]
 
     canonicalQuery = queries.canonicalEvents(
+        awc,
         queries.DesktopQueryParams(
             bid_window="aw-watcher-window_",
             bid_afk="aw-watcher-afk_",
             classes=categories,
-        )
+        ),
     )
     res = awc.query(
         f"""

examples/working_hours.py

@@ -3,7 +3,6 @@
 
 Also saves the matching work-events to a JSON file (for auditing purposes).
 """
-
 import json
 import logging
 import os
@@ -69,12 +68,13 @@ def query(regex: str, timeperiods, hostname: str):
     aw = aw_client.ActivityWatchClient()
 
     canonicalQuery = queries.canonicalEvents(
+        aw,
         queries.DesktopQueryParams(
             bid_window=f"aw-watcher-window_{hostname}",
             bid_afk=f"aw-watcher-afk_{hostname}",
             classes=categories,
             filter_classes=[["Work"]],
-        )
+        ),
     )
     query = f"""
     {canonicalQuery}