feat: default to using server-side categories for queries

ErikBjare · ErikBjare · commit 75ae5912afbe · 2023-11-22T08:51:04.000+01:00
diff --git a/aw_client/classes.py b/aw_client/classes.py
@@ -3,8 +3,18 @@
 
 Taken from default classes in aw-webui
 """
+import logging
+import random
+from typing import (
+    Any,
+    Dict,
+    List,
+    Tuple,
+)
 
-from typing import List, Dict, Any, Tuple
+import aw_client
+
+logger = logging.getLogger(__name__)
 
 CategoryId = List[str]
 CategorySpec = Dict[str, Any]
@@ -52,3 +62,20 @@
         {"type": "regex", "regex": "Gmail|Thunderbird|mutt|alpine"},
     ),
 ]
+
+
+def get_classes() -> 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.
+    """
+    awc = aw_client.ActivityWatchClient(f"get-setting-{random.randint(0, 10000)}")
+    try:
+        classes = awc.get_setting("classes")
+    except Exception:
+        logger.warning(
+            "Failed to get classes from server, using default classes as fallback"
+        )
+        return default_classes
+    # map into list of tuples
+    return [(v["name"], v["rule"]) for v in classes]
diff --git a/aw_client/client.py b/aw_client/client.py
@@ -341,12 +341,11 @@ def query(
     # Settings
     #
 
-    def get_setting(self, key=None) -> dict:
-        # TODO: explicitly fetch key from server, instead of fetching all settings
-        settings = self._get("settings").json()
+    def get_setting(self, key: Optional[str] = None) -> dict:
         if key:
-            return settings.get(key, None)
-        return settings
+            return self._get(f"settings/{key}").json()
+        else:
+            return self._get("settings").json()
 
     def set_setting(self, key: str, value: str) -> None:
         self._post(f"settings/{key}", value)
diff --git a/aw_client/queries.py b/aw_client/queries.py
@@ -3,18 +3,24 @@
 
 Most of these are from: https://github.com/ActivityWatch/aw-webui/blob/master/src/queries.ts
 """
-
+import dataclasses
 import json
 import re
-import dataclasses
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
+from typing import (
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
 
-from typing import List, Union, Tuple, Optional
 from typing_extensions import TypeGuard
 
 import aw_client
 
+from .classes import get_classes
+
 
 class EnhancedJSONEncoder(json.JSONEncoder):
     """For encoding dataclasses into JSON"""
@@ -75,6 +81,11 @@ def isAndroidParams(params: QueryParams) -> TypeGuard[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)
