flake/scripts/spectacle-screenshot.py

#! /usr/bin/env nix-shell
#! nix-shell -i python -p python312

import argparse
import os
import subprocess
import tempfile
from pathlib import Path
from shutil import which

from common.common import notify


def spectacle_screenshot(
    url: str | None = None, record: bool = False, file_path: Path | None = None
) -> None:
    try:
        if not which("spectacle"):
            raise FileNotFoundError("spectacle is not installed.")

        if file_path and Path(file_path).exists():
            raise FileExistsError(
                'File already exists. Please provide a different file path, or use the zipline function to upload the file.\nExample: zipline "{file_path}"'
            )

        if not file_path:
            # Spectacle actually defaults to .webm for video recordings,
            # but Zipline doesn't behave well with .webm files so we use .mp4 instead.
            use_temp_file = True
            suffix = ".mp4" if record else ".png"
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            file_path = temp_file.name
            temp_file.close()
        else:
            use_temp_file = False

        command = [
            "spectacle",
            "--nonotify",
            "--background",
            "--pointer",
            "--copy-image",
            "--output",
            file_path,
        ]

        if record:
            command.append("--record=region")
        else:
            command.append("--region")

        try:
            subprocess.run(command, check=True)
        except subprocess.CalledProcessError as e:
            if Path(file_path).exists() and use_temp_file:
                os.remove(file_path)
            raise e

        if not Path(file_path).stat().st_size:
            os.remove(file_path)
            raise FileNotFoundError("The file was not created properly.")

        try:
            opts = [
                "/etc/nixos/scripts/zipline.py",
                file_path,
                "--application-name",
                "Spectacle",
                "--desktop-entry",
                "org.kde.spectacle",
            ]
            if url:
                opts.extend(["--url", url])
            subprocess.run(opts)
        finally:
            if Path(file_path).exists() and use_temp_file:
                os.remove(file_path)
    except Exception as e:
        notify(
            application_name="Spectacle",
            title="An error occurred",
            message=str(e),
            urgency="critical",
            category="transfer.error",
            desktop_entry="org.kde.spectacle",
        )
        raise e


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="spectacle-screenshot.py",
        description="Take a screenshot or recording with Spectacle and automatically upload it to a Zipline instance.",
        epilog="Example usage: spectacle-screenshot.py",
    )
    parser.add_argument(
        "--url",
        help="The URL of the Zipline instance to upload the screenshot or recording to. Defaults to 'https://csw.im'.",
        default="https://csw.im",
    )
    parser.add_argument(
        "--record",
        help="If this is set, Spectacle will record the region instead of taking a screenshot.",
        action="store_true",
    )
    parser.add_argument(
        "--file-path",
        help="The path to save the screenshot or recording to. If not provided, the screenshot or recording will be saved to a temporary file.",
        default=None,
    )

    args = parser.parse_args()

    spectacle_screenshot(
        url=args.url,
        record=args.record,
        file_path=args.file_path,
    )