upstream-aw-watcher-window/aw_watcher_window/main.py

import logging
import os
import re
import signal
import subprocess
import sys
from datetime import datetime, timezone
from time import sleep

from aw_client import ActivityWatchClient
from aw_core.log import setup_logging
from aw_core.models import Event

from .config import parse_args
from .exceptions import FatalError
from .lib import get_current_window
from .research_filter import transform as research_transform
from .macos_cli import build_swift_command
from .macos_permissions import background_ensure_permissions

logger = logging.getLogger(__name__)

# run with LOG_LEVEL=DEBUG
log_level = os.environ.get("LOG_LEVEL")
if log_level:
    logger.setLevel(logging.__getattribute__(log_level.upper()))


def compute_pulsetime(poll_time: float) -> float:
    """Scale pulsetime with poll_time so OS scheduling jitter doesn't break heartbeat chains.

    At poll_time=1s, jitter ~0.15s is well within 1s margin (poll_time+1).
    At poll_time=5s, jitter ~0.75s exceeds the 1s margin ~10% of the time,
    causing missing time in the timeline. max(poll_time*1.5, poll_time+1) keeps
    backward compatibility at poll_time≤2s while fixing the problem at higher
    polling intervals. See: https://github.com/ActivityWatch/activitywatch/issues/1177
    """
    return max(poll_time * 1.5, poll_time + 1.0)


def kill_process(pid):
    logger.info("Killing process {}".format(pid))
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        logger.info("Process {} already dead".format(pid))


def try_compile_title_regex(title):
    try:
        return re.compile(title, re.IGNORECASE)
    except re.error:
        logger.error(f"Invalid regex pattern: {title}")
        exit(1)


def main():
    args = parse_args()

    if sys.platform.startswith("linux") and (
        "DISPLAY" not in os.environ or not os.environ["DISPLAY"]
    ):
        raise Exception("DISPLAY environment variable not set")

    setup_logging(
        name="aw-watcher-window",
        testing=args.testing,
        verbose=args.verbose,
        log_stderr=True,
        log_file=True,
    )
    if sys.platform == "darwin":
        background_ensure_permissions()

    client = ActivityWatchClient(
        "aw-watcher-window", host=args.host, port=args.port, testing=args.testing
    )

    bucket_id = f"{client.client_name}_{client.client_hostname}"
    event_type = "currentwindow"

    client.create_bucket(bucket_id, event_type, queued=True)

    logger.info("aw-watcher-window started")
    client.wait_for_start()

    with client:
        research_category_map = (
            args.research_category_map
            if args.research_enabled
            else None
        )
        research_app_category_map = (
            args.research_app_category_map
            if args.research_enabled
            else None
        )
        if sys.platform == "darwin" and args.strategy == "swift":
            logger.info("Using swift strategy, calling out to swift binary")
            binpath = os.path.join(
                os.path.dirname(os.path.realpath(__file__)), "aw-watcher-window-macos"
            )

            try:
                p = subprocess.Popen(
                    build_swift_command(
                        binpath,
                        client.server_address,
                        bucket_id,
                        client.client_hostname,
                        client.client_name,
                        exclude_title=args.exclude_title,
                        exclude_titles=args.exclude_titles,
                        research_category_map=research_category_map,
                        research_app_category_map=research_app_category_map,
                    )
                )
                # terminate swift process when this process dies
                signal.signal(signal.SIGTERM, lambda *_: kill_process(p.pid))
                p.wait()
            except KeyboardInterrupt:
                print("KeyboardInterrupt")
                kill_process(p.pid)
        else:
            heartbeat_loop(
                client,
                bucket_id,
                poll_time=args.poll_time,
                strategy=args.strategy,
                exclude_title=args.exclude_title,
                exclude_titles=[
                    try_compile_title_regex(title)
                    for title in args.exclude_titles
                    if title is not None
                ],
                research_category_map=research_category_map,
                research_app_category_map=research_app_category_map,
            )


def heartbeat_loop(
    client,
    bucket_id,
    poll_time,
    strategy,
    exclude_title=False,
    exclude_titles=[],
    research_category_map=None,
    research_app_category_map=None,
):
    while True:
        if os.getppid() == 1:
            logger.info("window-watcher stopped because parent process died")
            break

        current_window = None
        try:
            current_window = get_current_window(strategy)
            logger.debug(current_window)
        except (FatalError, OSError):
            # Fatal exceptions should quit the program
            try:
                logger.exception("Fatal error, stopping")
            except OSError:
                pass
            break
        except Exception:
            # Non-fatal exceptions should be logged
            try:
                # If stdout has been closed, this exception-print can cause (I think)
                #   OSError: [Errno 5] Input/output error
                # See: https://github.com/ActivityWatch/activitywatch/issues/756#issue-1296352264
                #
                # However, I'm unable to reproduce the OSError in a test (where I close stdout before logging),
                # so I'm in uncharted waters here... but this solution should work.
                logger.exception("Exception thrown while trying to get active window")
            except OSError:
                break

        if current_window is None:
            logger.debug("Unable to fetch window, trying again on next poll")
        else:
            current_window = transform_window(
                current_window,
                exclude_title=exclude_title,
                exclude_titles=exclude_titles,
                research_category_map=research_category_map,
                research_app_category_map=research_app_category_map,
            )

            now = datetime.now(timezone.utc)
            current_window_event = Event(timestamp=now, data=current_window)

            client.heartbeat(
                bucket_id,
                current_window_event,
                pulsetime=compute_pulsetime(poll_time),
                queued=True,
            )

        sleep(poll_time)


def transform_window(
    current_window,
    exclude_title=False,
    exclude_titles=None,
    research_category_map=None,
    research_app_category_map=None,
):
    if research_category_map is not None:
        return research_transform(
            current_window,
            research_category_map,
            app_category_map=research_app_category_map,
        )

    for pattern in exclude_titles or []:
        if pattern.search(current_window["title"]):
            current_window["title"] = "excluded"

    if exclude_title:
        current_window["title"] = "excluded"

    return current_window