flake/scripts/py/spectacle-screenshot.py

#! /usr/bin/env nix-shell
#! nix-shell /etc/nixos/scripts/nix/python.nix -i python

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

from common.common import notify, run
from PySide6.QtGui import QImageWriter  # type: ignore


def spectacle_screenshot(
    url: str | None = None,
    record: bool = False,
    file_path: Path | None = None,
    format: str | 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}"'
            )

        accepted_formats = [
            f.data().decode() for f in QImageWriter.supportedImageFormats()
        ]
        accepted_formats.extend(("webm", "mp4"))
        print(accepted_formats)
        if not format:
            format = "mp4" if record else "png"
        if format.lower() not in accepted_formats:
            raise ValueError(
                f"Invalid format. Accepted formats are: {', '.join(accepted_formats)}"
            )
        else:
            format = f".{str(format)}"

        if not file_path:
            use_temp_file = True
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=format)
            file_path = Path(temp_file.name)
            temp_file.close()
        else:
            use_temp_file = False

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

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

        try:
            run(command)
        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/py/zipline.py",
                str(file_path),
                "--application-name",
                "Spectacle",
                "--desktop-entry",
                "org.kde.spectacle",
            ]
            if url:
                opts.extend(["--url", url])
            run(cmd=opts)  # type: ignore
        finally:
            if Path(file_path).exists() and use_temp_file:
                os.remove(file_path)
    except BaseException as e:
        notify(
            application_name="Spectacle",
            title="An error occurred",
            message=str(e),
            urgency="normal",
            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(
        "--format",
        help="The format of the screenshot. Defaults to 'png' if `--record` is not set, and 'mp4' if `--record` is set.",
        default=None,
    )
    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,
        format=args.format,
    )