feat: rewrote CLI to use click

Author: ErikBjare

aw_client/cli.py

@@ -4,9 +4,15 @@
 from datetime import timedelta, datetime, timezone
 from pprint import pprint
 
+import click
+
 import aw_client
 from aw_core import Event
 
+now = datetime.now(timezone.utc)
+td1day = timedelta(days=1)
+td1yr = timedelta(days=365)
+
 
 def _valid_date(s):
     # https://stackoverflow.com/questions/25470844/specify-format-for-input-arguments-argparse-python
@@ -17,110 +23,104 @@ def _valid_date(s):
         raise argparse.ArgumentTypeError(msg)
 
 
-def main():
-    now = datetime.now(timezone.utc)
-    td1day = timedelta(days=1)
-    td1yr = timedelta(days=365)
-
-    parser = argparse.ArgumentParser(
-        prog="aw-cli", description="A CLI utility for interacting with ActivityWatch."
-    )
-    parser.set_defaults(which="none")
-    parser.add_argument(
-        "--host",
-        default="127.0.0.1:5600",
-        help="Host to use, in the format HOSTNAME[:PORT]",
-    )
-    parser.add_argument(
-        "--testing", action="store_true", help="Set to use testing ports by default"
-    )
-
-    subparsers = parser.add_subparsers(help="sub-command help")
-
-    parser_heartbeat = subparsers.add_parser(
-        "heartbeat", help="Send a heartbeat to the server"
-    )
-    parser_heartbeat.set_defaults(which="heartbeat")
-    parser_heartbeat.add_argument("--pulsetime", default=60, help="Pulsetime to use")
-    parser_heartbeat.add_argument("bucket", help="bucketname to send heartbeat to")
-    parser_heartbeat.add_argument(
-        "data", default="{}", help="JSON data to send in heartbeat"
-    )
-
-    parser_buckets = subparsers.add_parser("buckets", help="List all buckets")
-    parser_buckets.set_defaults(which="buckets")
-
-    parser_buckets = subparsers.add_parser("events", help="Query events from bucket")
-    parser_buckets.set_defaults(which="events")
-    parser_buckets.add_argument("bucket")
-
-    parser_query = subparsers.add_parser("query", help="Query events from bucket")
-    parser_query.set_defaults(which="query")
-    parser_query.add_argument("path")
-    parser_query.add_argument("--name")
-    parser_query.add_argument("--cache", action="store_true")
-    parser_query.add_argument(
-        "--json", action="store_true", help="Output resulting JSON"
+@click.group(
+    help="CLI utility for aw-client to aid in interacting with the ActivityWatch server"
+)
+@click.option("--testing", is_flag=True, help="Set to use testing ports by default")
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="Address of host",
+)
+@click.option(
+    "--port",
+    default=5600,
+    help="Port to use",
+)
+@click.pass_context
+def main(ctx, testing: bool, host: str, port: int):
+    ctx.testing = testing
+    ctx.client = aw_client.ActivityWatchClient(
+        host=host,
+        port=port if port != 5600 else (5666 if testing else 5600),
     )
-    parser_query.add_argument("--start", default=now - td1day, type=_valid_date)
-    parser_query.add_argument("--end", default=now + 10 * td1yr, type=_valid_date)
 
-    args = parser.parse_args()
-    # print("Args: {}".format(args))
 
-    client = aw_client.ActivityWatchClient(
-        host=args.host.split(":")[0],
-        port=int(
-            (args.host.split(":")[1:] + [5600 if not args.testing else 5666]).pop()
-        ),
-    )
-
-    if args.which == "heartbeat":
-        e = Event(duration=0, data=json.loads(args.data), timestamp=now)
-        print(e)
-        client.heartbeat(args.bucket, e, args.pulsetime)
-    elif args.which == "buckets":
-        buckets = client.get_buckets()
-        print("Buckets:")
-        for bucket in buckets:
-            print(" - {}".format(bucket))
-    elif args.which == "events":
-        events = client.get_events(args.bucket)
-        print("events:")
-        for e in events:
-            print(
-                " - {} ({}) {}".format(
-                    e.timestamp.replace(tzinfo=None, microsecond=0),
-                    str(e.duration).split(".")[0],
-                    e.data,
-                )
+@main.command(help="Send a heartbeat to bucket with ID `bucket_id` with JSON `data`")
+@click.argument("bucket_id")
+@click.argument("data")
+@click.option("--pulsetime", default=60, help="pulsetime to use for merging heartbeats")
+@click.pass_context
+def heartbeat(ctx, bucket_id: str, data: str, pulsetime: int):
+    now = datetime.now(timezone.utc)
+    e = Event(duration=0, data=json.loads(data), timestamp=now)
+    print(e)
+    ctx.client.heartbeat(bucket_id, e, pulsetime)
+
+
+@main.command(help="List all buckets")
+@click.pass_context
+def buckets(ctx):
+    buckets = ctx.client.get_buckets()
+    print("Buckets:")
+    for bucket in buckets:
+        print(" - {}".format(bucket))
+
+
+@main.command(help="Query events from bucket with ID `bucket_id`")
+@click.argument("bucket_id")
+@click.pass_context
+def events(ctx, bucket_id: str):
+    events = ctx.client.get_events(bucket_id)
+    print("events:")
+    for e in events:
+        print(
+            " - {} ({}) {}".format(
+                e.timestamp.replace(tzinfo=None, microsecond=0),
+                str(e.duration).split(".")[0],
+                e.data,
             )
-    elif args.which == "query":
-        with open(args.path) as f:
-            query = f.read()
-        result = client.query(
-            query, args.start, args.end, cache=args.cache, name=args.name
         )
-        if args.json:
-            print(json.dumps(result))
-        else:
-            for period in result:
-                print("Showing 10 out of {} events:".format(len(period)))
-                for event in period[:10]:
-                    event.pop("id")
-                    event.pop("timestamp")
-                    print(
-                        " - Duration: {} \tData: {}".format(
-                            str(timedelta(seconds=event["duration"])).split(".")[0],
-                            event["data"],
-                        )
-                    )
+
+
+@main.command(help="Run a query in file at `path` on the server")
+@click.argument("path")
+@click.option("--name")
+@click.option("--cache", is_flag=True)
+@click.option("--json", is_flag=True)
+@click.option("--start", default=now - td1day, type=click.DateTime())
+@click.option("--stop", default=now + td1yr, type=click.DateTime())
+@click.pass_context
+def query(
+    ctx,
+    path: str,
+    cache: bool,
+    _json: bool,
+    start: datetime,
+    stop: datetime,
+    name: str = None,
+):
+    with open(path) as f:
+        query = f.read()
+    result = ctx.client.query(query, [(start, stop)], cache=cache, name=name)
+    if _json:
+        print(json.dumps(result))
+    else:
+        for period in result:
+            print("Showing 10 out of {} events:".format(len(period)))
+            for event in period[:10]:
+                event.pop("id")
+                event.pop("timestamp")
                 print(
-                    "Total duration:\t",
-                    timedelta(seconds=sum(e["duration"] for e in period)),
+                    " - Duration: {} \tData: {}".format(
+                        str(timedelta(seconds=event["duration"])).split(".")[0],
+                        event["data"],
+                    )
                 )
-    else:
-        parser.print_help()
+            print(
+                "Total duration:\t",
+                timedelta(seconds=sum(e["duration"] for e in period)),
+            )
 
 
 if __name__ == "__main__":

aw_client/client.py

@@ -6,7 +6,7 @@
 import functools
 from datetime import datetime
 from collections import namedtuple
-from typing import Optional, List, Any, Union, Dict, Callable
+from typing import Optional, List, Any, Union, Dict, Callable, Tuple
 
 import requests as req
 import persistqueue
@@ -287,11 +287,10 @@ def import_bucket(self, bucket: dict) -> None:
     def query(
         self,
         query: str,
-        start: datetime,
-        end: datetime,
+        timeperiods: List[Tuple[datetime, datetime]],
         name: str = None,
         cache: bool = False,
-    ) -> Union[int, dict]:
+    ) -> Any:
         endpoint = "query/"
         params = {}  # type: Dict[str, Any]
         if cache:
@@ -301,8 +300,12 @@ def query(
                 )
             params["name"] = name
             params["cache"] = int(cache)
+
         data = {
-            "timeperiods": ["/".join([start.isoformat(), end.isoformat()])],
+            "timeperiods": [
+                "/".join([start.isoformat(), end.isoformat()])
+                for start, end in timeperiods
+            ],
             "query": query.split("\n"),
         }
         response = self._post(endpoint, data, params=params)