cleaned up main function and started using the client context manager

aw_watcher_window/main.py

@@ -22,33 +22,40 @@ def main():
     #   req_version is 3.5 due to usage of subprocess.run
     assert_version((3, 5))
 
+    if sys.platform.startswith("linux") and ("DISPLAY" not in os.environ or not os.environ["DISPLAY"]):
+        raise Exception("DISPLAY environment variable not set")
+
     # Read settings from config
     config = load_config()
-
-    # Parse arguments
-    parser = argparse.ArgumentParser("A cross platform window watcher for Linux, macOS and Windows.")
-    parser.add_argument("--testing", dest="testing", action="store_true")
-    parser.add_argument("--exclude-title", dest="exclude_title", action="store_true")
-    parser.add_argument("--verbose", dest="verbose", action="store_true")
-    parser.add_argument("--poll-time", type=float, default=config.getfloat("poll_time"))
-    args = parser.parse_args()
+    args = parse_args(default_poll_time=config.getfloat("poll_time"))
 
     setup_logging(name="aw-watcher-window", testing=args.testing, verbose=args.verbose,
                   log_stderr=True, log_file=True)
 
-    logging.info("Running watcher with poll time {} seconds".format(args.poll_time))
-
-    if sys.platform.startswith("linux") and ("DISPLAY" not in os.environ or not os.environ["DISPLAY"]):
-        raise Exception("DISPLAY environment variable not set")
+    logger.info("Running watcher with poll time {} seconds".format(args.poll_time))
 
     client = ActivityWatchClient("aw-watcher-window", testing=args.testing)
 
-    bucketname = "{}_{}".format(client.client_name, client.client_hostname)
-    eventtype = "currentwindow"
-    client.setup_bucket(bucketname, eventtype)
-    client.connect()
+    bucket_id = "{}_{}".format(client.client_name, client.client_hostname)
+    event_type = "currentwindow"
+    client.create_bucket(bucket_id, event_type, queued=True)
 
     logger.info("aw-watcher-window has started")
+    with client:
+        heartbeat_loop(client, bucket_id, poll_time=args.poll_time, exclude_title=args.exclude_title)
+
+
+def parse_args(default_poll_time: float):
+    """config contains defaults loaded from the config file"""
+    parser = argparse.ArgumentParser("A cross platform window watcher for Activitywatch.\nSupported on: Linux (X11), macOS and Windows.")
+    parser.add_argument("--testing", dest="testing", action="store_true")
+    parser.add_argument("--exclude-title", dest="exclude_title", action="store_true")
+    parser.add_argument("--verbose", dest="verbose", action="store_true")
+    parser.add_argument("--poll-time", dest="poll_time", type=float, default=default_poll_time)
+    return parser.parse_args()
+
+
+def heartbeat_loop(client, bucket_id, poll_time, exclude_title=False):
     while True:
         try:
             current_window = get_current_window()
@@ -65,14 +72,14 @@ def main():
             # Create current_window event
             data = {
                 "app": current_window["appname"],
-                "title": current_window["title"] if not args.exclude_title else "title:excluded"
+                "title": current_window["title"] if not exclude_title else "excluded"
             }
             current_window_event = Event(timestamp=now, data=data)
 
             # Set pulsetime to 1 second more than the poll_time
             # This since the loop takes more time than poll_time
             # due to sleep(poll_time).
-            client.heartbeat(bucketname, current_window_event,
-                             pulsetime=args.poll_time + 1.0, queued=True)
+            client.heartbeat(bucket_id, current_window_event,
+                             pulsetime=poll_time + 1.0, queued=True)
 
-        sleep(args.poll_time)
+        sleep(poll_time)