xhs-note-crawling/workplace/screenshot.py

from __future__ import annotations

from abc import ABC, abstractmethod
from pathlib import Path

import numpy as np
from PIL import ImageGrab

_REPOSITORY_ROOT_DIRECTORY = Path(__file__).resolve().parent.parent
_OUTPUT_SCREENSHOT_FILE_PATH = _REPOSITORY_ROOT_DIRECTORY / "output" / "screenshot.png"

FULL_SCREEN_SCREENSHOT_OUTPUT_PNG_FILE_PATH: Path = _OUTPUT_SCREENSHOT_FILE_PATH

current_full_screen_screenshot_bgr_numpy: np.ndarray | None = None


class FullScreenScreenshotCaptureSaver(ABC):
    @abstractmethod
    def capture_full_screen_and_store_in_memory(self) -> np.ndarray:
        ...


class PillowImageGrabFullScreenScreenshotCaptureSaver(FullScreenScreenshotCaptureSaver):
    def capture_full_screen_and_store_in_memory(self) -> np.ndarray:
        global current_full_screen_screenshot_bgr_numpy
        full_screen_rgba_image = ImageGrab.grab()
        full_screen_rgb_image = full_screen_rgba_image.convert("RGB")
        _OUTPUT_SCREENSHOT_FILE_PATH.parent.mkdir(parents=True, exist_ok=True)
        full_screen_rgb_image.save(_OUTPUT_SCREENSHOT_FILE_PATH, format="PNG")
        rgb_uint8_numpy = np.asarray(full_screen_rgb_image, dtype=np.uint8)
        bgr_uint8_numpy = rgb_uint8_numpy[:, :, ::-1].copy()
        current_full_screen_screenshot_bgr_numpy = bgr_uint8_numpy
        return bgr_uint8_numpy


def capture_full_screen_screenshot_bgr_numpy_and_save_screenshot_png() -> np.ndarray:
    """全屏截图：写入 ``output/screenshot.png``（覆盖），并返回 BGR ``numpy``（与内存缓存一致）。"""
    full_screen_screenshot_capture_saver = PillowImageGrabFullScreenScreenshotCaptureSaver()
    return full_screen_screenshot_capture_saver.capture_full_screen_and_store_in_memory()