553 lines
20 KiB
Python
553 lines
20 KiB
Python
# This is part of the MSS Python's module.
|
|
# Source: https://github.com/BoboTiG/python-mss.
|
|
|
|
from __future__ import annotations
|
|
|
|
import platform
|
|
import warnings
|
|
from abc import ABC, abstractmethod
|
|
from datetime import datetime
|
|
from threading import Lock
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from mss.exception import ScreenShotError
|
|
from mss.screenshot import ScreenShot
|
|
from mss.tools import to_png
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable, Iterator
|
|
|
|
from mss.models import Monitor, Monitors, Size
|
|
|
|
# Prior to 3.11, Python didn't have the Self type. typing_extensions does, but we don't want to depend on it.
|
|
try:
|
|
from typing import Self
|
|
except ImportError:
|
|
try:
|
|
from typing_extensions import Self
|
|
except ImportError:
|
|
Self = Any # type: ignore[assignment]
|
|
|
|
try:
|
|
from datetime import UTC
|
|
except ImportError:
|
|
# Python < 3.11
|
|
from datetime import timezone
|
|
|
|
UTC = timezone.utc
|
|
|
|
#: Global lock protecting access to platform screenshot calls.
|
|
#:
|
|
#: .. versionadded:: 6.0.0
|
|
#:
|
|
#: .. deprecated:: 10.2.0
|
|
#: The global lock is no longer used, and will be removed in a future release.
|
|
#: MSS objects now have their own locks, which are not publicly-accessible.
|
|
lock = Lock()
|
|
|
|
OPAQUE = 255
|
|
|
|
|
|
# A sentinel value to indicate that a parameter was not passed, as opposed to being passed with a value of None. This
|
|
# is used in the MSS constructor to distinguish between the user not passing a parameter, and the user explicitly
|
|
# passing None (which is the default for some parameters). This allows us to preserve the existing behavior of ignoring
|
|
# certain parameters on certain platforms, while still allowing users to explicitly set those parameters on platforms
|
|
# where they are supported.
|
|
class _PlatformSpecific:
|
|
def __init__(self, sphinx_repr: Any) -> None:
|
|
self.sphinx_repr = str(sphinx_repr)
|
|
|
|
def __repr__(self) -> str:
|
|
# This is used to get Sphinx to show a useful default when it shows the default in the summary, rather than
|
|
# an opaque object.
|
|
return self.sphinx_repr
|
|
|
|
|
|
__all__ = ()
|
|
|
|
|
|
class MSSImplementation(ABC):
|
|
"""Base class for internal platform/backend implementations.
|
|
|
|
Only one of these methods will be called at a time; the containing
|
|
MSS object will hold a lock during these calls.
|
|
"""
|
|
|
|
__slots__ = ("performance_status", "with_cursor")
|
|
|
|
with_cursor: bool
|
|
|
|
def __init__(self, /, *, with_cursor: bool = False) -> None:
|
|
# We put with_cursor on the MSSImplementation because the Xlib backend will turn it off if the library isn't
|
|
# installed. (It's not a separate library under XCB.) So, we need to let the backend mutate it. Note that
|
|
# the other platforms don't support with_cursor, and don't pass it to us.
|
|
#
|
|
# TODO(jholveck): #493 We should remove this expectation in 11.0. It seems unlikely to be practically useful,
|
|
# Xlib is legacy, and just complicates things.
|
|
self.with_cursor = with_cursor
|
|
|
|
# Any notes the backend needs to give the user for debugging purposes, like why it had to fall back to a
|
|
# slower implementation.
|
|
self.performance_status: list[str] = []
|
|
|
|
@abstractmethod
|
|
def cursor(self) -> ScreenShot | None:
|
|
"""Retrieve all cursor data. Pixels have to be RGB."""
|
|
|
|
@abstractmethod
|
|
def grab(self, monitor: Monitor, /) -> bytearray | tuple[bytearray, Size]:
|
|
"""Retrieve all pixels from a monitor. Pixels have to be RGB.
|
|
|
|
If the monitor size is not in pixel units, include a Size in
|
|
pixels (see issue #23).
|
|
"""
|
|
|
|
@abstractmethod
|
|
def monitors(self) -> Monitors:
|
|
"""Return positions of monitors."""
|
|
|
|
def close(self) -> None: # noqa: B027 - intentionally empty
|
|
"""Clean up.
|
|
|
|
This will be called at most once.
|
|
|
|
It's not necessary for subclasses to implement this if they
|
|
have nothing to clean up.
|
|
"""
|
|
|
|
@staticmethod
|
|
def _cfactory(
|
|
attr: Any,
|
|
func: str,
|
|
argtypes: list[Any],
|
|
restype: Any,
|
|
/,
|
|
errcheck: Callable | None = None,
|
|
) -> None:
|
|
"""Factory to create a ctypes function and automatically manage errors."""
|
|
meth = getattr(attr, func)
|
|
meth.argtypes = argtypes
|
|
meth.restype = restype
|
|
if errcheck:
|
|
meth.errcheck = errcheck
|
|
|
|
|
|
def _choose_impl(**kwargs: Any) -> MSSImplementation:
|
|
"""Return the backend implementation for the current platform.
|
|
|
|
Detects the platform we are running on and instantiates the
|
|
appropriate internal implementation class.
|
|
|
|
.. seealso::
|
|
- :class:`mss.MSS`
|
|
- :class:`mss.darwin.MSS`
|
|
- :class:`mss.linux.MSS`
|
|
- :class:`mss.windows.MSS`
|
|
"""
|
|
os_ = platform.system().lower()
|
|
|
|
if os_ == "darwin":
|
|
from mss.darwin import MSSImplDarwin # noqa: PLC0415
|
|
|
|
return MSSImplDarwin(**kwargs)
|
|
|
|
if os_ == "linux":
|
|
from mss.linux import choose_impl as choose_impl_linux # noqa: PLC0415
|
|
|
|
# Linux has its own factory to choose the backend.
|
|
return choose_impl_linux(**kwargs)
|
|
|
|
if os_ == "windows":
|
|
from mss.windows import choose_impl as choose_impl_windows # noqa: PLC0415
|
|
|
|
# Windows has its own factory to choose the backend.
|
|
return choose_impl_windows(**kwargs)
|
|
|
|
msg = f"System {os_!r} not (yet?) implemented."
|
|
raise ScreenShotError(msg)
|
|
|
|
|
|
# Does this belong here?
|
|
class MSS:
|
|
"""Multiple ScreenShots class
|
|
|
|
:param backend: Backend selector, for platforms with multiple backends.
|
|
:param compression_level: PNG compression level.
|
|
:param with_cursor: Include the mouse cursor in screenshots
|
|
(GNU/Linux only)
|
|
:type display: bool, optional (default False)
|
|
:param display: X11 display name (GNU/Linux only).
|
|
:type display: bytes | str, optional (default :envvar:`$DISPLAY`)
|
|
:param max_displays: Maximum number of displays to enumerate (macOS only).
|
|
:type max_displays: int, optional (default 32)
|
|
|
|
.. versionadded:: 8.0.0
|
|
``compression_level``, ``display``, ``max_displays``, and
|
|
``with_cursor`` keyword arguments.
|
|
|
|
.. versionadded:: 10.2.0
|
|
``backend`` keyword argument.
|
|
"""
|
|
|
|
# We want to:
|
|
# * Let Sphinx, IDEs, and code-checkers know all the possible kwargs.
|
|
# * Know if a user explicitly passed an unsupported platform-dependent keyword, so we can warn.
|
|
# * Show a meaningful default in the Sphinx doc's summary string
|
|
#
|
|
# To accomplish this:
|
|
# * We list the possibilities explicitly in the __init__ kwargs.
|
|
# * We use a sentinel value, so we can tell whether or not the user actually gave us a value.
|
|
# * We represent the "default value" sentinel object with something different, so Sphinx formats it usefully.
|
|
_PD_WITH_CURSOR = _PlatformSpecific(False) # noqa: FBT003
|
|
_PD_DISPLAY = _PlatformSpecific(None)
|
|
_PD_MAX_DISPLAYS = _PlatformSpecific(32)
|
|
|
|
def __init__(
|
|
self,
|
|
/,
|
|
*,
|
|
backend: str = "default",
|
|
compression_level: int = 6,
|
|
with_cursor: bool | _PlatformSpecific = _PD_WITH_CURSOR,
|
|
display: bytes | str | None | _PlatformSpecific = _PD_DISPLAY,
|
|
max_displays: int | _PlatformSpecific = _PD_MAX_DISPLAYS,
|
|
) -> None:
|
|
impl_kwargs = {}
|
|
|
|
system = platform.system()
|
|
for name, value, supported_platforms in [
|
|
("with_cursor", with_cursor, ["Linux"]),
|
|
("display", display, ["Linux"]),
|
|
("max_displays", max_displays, ["Darwin"]),
|
|
]:
|
|
if isinstance(value, _PlatformSpecific):
|
|
continue
|
|
if system not in supported_platforms:
|
|
# TODO(jholveck): #493 Accept platform-specific kwargs on all platforms for migration ease. Foreign
|
|
# kwargs are silently stripped with a warning.
|
|
warnings.warn(
|
|
f"{name} is only available on {', '.join(supported_platforms)}. "
|
|
"This will be an error in the future.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
else:
|
|
impl_kwargs[name] = value
|
|
|
|
self._impl: MSSImplementation = _choose_impl(
|
|
backend=backend,
|
|
**impl_kwargs,
|
|
)
|
|
|
|
# The cls_image is only used atomically, so does not require locking.
|
|
self.cls_image: type[ScreenShot] = ScreenShot
|
|
# The compression level is only used atomically, so does not require locking.
|
|
#: PNG compression level used when saving the screenshot data into a file
|
|
#: (see :py:func:`zlib.compress()` for details).
|
|
#:
|
|
#: .. versionadded:: 3.2.0
|
|
self.compression_level = compression_level
|
|
|
|
# The attributes below are protected by self._lock. The attributes above are user-visible, so we don't
|
|
# control when they're modified. Currently, we only make sure that they're safe to modify while locked, or
|
|
# document that the user shouldn't change them. We could also use properties protect them against changes, or
|
|
# change them under the lock.
|
|
self._lock = Lock()
|
|
self._monitors: Monitors | None = None
|
|
self._closed = False
|
|
|
|
def __enter__(self) -> Self:
|
|
"""For the cool call `with MSS() as mss:`."""
|
|
return self
|
|
|
|
def __exit__(self, *_: object) -> None:
|
|
"""For the cool call `with MSS() as mss:`."""
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
"""Clean up.
|
|
|
|
This releases resources that MSS may be using. Once the MSS
|
|
object is closed, it may not be used again.
|
|
|
|
It is safe to call this multiple times; multiple calls have no
|
|
effect.
|
|
|
|
Rather than use :py:meth:`close` explicitly, we recommend you
|
|
use the MSS object as a context manager::
|
|
|
|
with mss.MSS() as sct:
|
|
...
|
|
"""
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
self._impl.close()
|
|
self._closed = True
|
|
|
|
def grab(self, monitor: Monitor | tuple[int, int, int, int], /) -> ScreenShot:
|
|
"""Retrieve screen pixels for a given monitor.
|
|
|
|
Note: ``monitor`` can be a tuple like the one
|
|
:py:meth:`PIL.ImageGrab.grab` accepts: ``(left, top, right, bottom)``
|
|
|
|
:param monitor: The coordinates and size of the box to capture.
|
|
See :meth:`monitors <monitors>` for object details.
|
|
:returns: Screenshot of the requested region.
|
|
"""
|
|
# Convert PIL bbox style
|
|
if isinstance(monitor, tuple):
|
|
monitor = {
|
|
"left": monitor[0],
|
|
"top": monitor[1],
|
|
"width": monitor[2] - monitor[0],
|
|
"height": monitor[3] - monitor[1],
|
|
}
|
|
|
|
if monitor["width"] <= 0 or monitor["height"] <= 0:
|
|
msg = f"Region has zero or negative size: {monitor!r}"
|
|
raise ScreenShotError(msg)
|
|
|
|
with self._lock:
|
|
img_data_and_maybe_size = self._impl.grab(monitor)
|
|
if isinstance(img_data_and_maybe_size, tuple):
|
|
img_data, size = img_data_and_maybe_size
|
|
screenshot = self.cls_image(img_data, monitor, size=size)
|
|
else:
|
|
img_data = img_data_and_maybe_size
|
|
screenshot = self.cls_image(img_data, monitor)
|
|
if self._impl.with_cursor and (cursor := self._impl.cursor()):
|
|
return self._merge(screenshot, cursor)
|
|
return screenshot
|
|
|
|
@property
|
|
def monitors(self) -> Monitors:
|
|
"""Get positions of all monitors.
|
|
If the monitor has rotation, you have to deal with it
|
|
inside this method.
|
|
|
|
This method has to fill ``self._monitors`` with all information
|
|
and use it as a cache:
|
|
|
|
- ``self._monitors[0]`` is a dict of all monitors together
|
|
- ``self._monitors[N]`` is a dict of the monitor N (with N > 0)
|
|
|
|
Each monitor is a dict with:
|
|
|
|
- ``left``: the x-coordinate of the upper-left corner
|
|
- ``top``: the y-coordinate of the upper-left corner
|
|
- ``width``: the width
|
|
- ``height``: the height
|
|
- ``is_primary``: (optional) true if this is the primary monitor
|
|
- ``name``: (optional) human-readable device name
|
|
- ``unique_id``: (optional) platform-specific stable identifier for the monitor
|
|
- ``output``: (optional, Linux only) monitor output name, compatible with xrandr
|
|
"""
|
|
with self._lock:
|
|
if self._monitors is None:
|
|
self._monitors = self._impl.monitors()
|
|
assert self._monitors is not None # noqa: S101
|
|
return self._monitors
|
|
|
|
@property
|
|
def primary_monitor(self) -> Monitor:
|
|
"""Get the primary monitor.
|
|
|
|
Returns the monitor marked as primary. If no monitor is marked as primary
|
|
(or the platform doesn't support primary monitor detection), returns the
|
|
first monitor (at index 1).
|
|
|
|
:raises ScreenShotError: If no monitors are available.
|
|
|
|
.. versionadded:: 10.2.0
|
|
"""
|
|
monitors = self.monitors
|
|
if len(monitors) <= 1: # Only the "all monitors" entry or empty
|
|
msg = "No monitor found."
|
|
raise ScreenShotError(msg)
|
|
|
|
return next(
|
|
(
|
|
monitor
|
|
for monitor in monitors[1:] # Skip the "all monitors" entry at index 0
|
|
if monitor.get("is_primary", False)
|
|
),
|
|
monitors[1], # Fallback to the first monitor if no primary is found
|
|
)
|
|
|
|
def save(
|
|
self,
|
|
/,
|
|
*,
|
|
mon: int = 0,
|
|
output: str = "monitor-{mon}.png",
|
|
callback: Callable[[str], None] | None = None,
|
|
) -> Iterator[str]:
|
|
"""Grab a screenshot and save it to a file.
|
|
|
|
:param int mon: The monitor to screenshot (default=0). ``-1`` grabs all
|
|
monitors, ``0`` grabs each monitor, and ``N`` grabs monitor ``N``.
|
|
:param str output: The output filename. Keywords: ``{mon}``, ``{top}``,
|
|
``{left}``, ``{width}``, ``{height}``, ``{date}``.
|
|
:param callable callback: Called before saving the screenshot; receives
|
|
the ``output`` argument.
|
|
:return: Created file(s).
|
|
"""
|
|
monitors = self.monitors
|
|
if not monitors:
|
|
msg = "No monitor found."
|
|
raise ScreenShotError(msg)
|
|
|
|
if mon == 0:
|
|
# One screenshot by monitor
|
|
for idx, monitor in enumerate(monitors[1:], 1):
|
|
fname = output.format(mon=idx, date=datetime.now(UTC) if "{date" in output else None, **monitor)
|
|
if callable(callback):
|
|
callback(fname)
|
|
sct = self.grab(monitor)
|
|
to_png(sct.rgb, sct.size, level=self.compression_level, output=fname)
|
|
yield fname
|
|
else:
|
|
# A screenshot of all monitors together or
|
|
# a screenshot of the monitor N.
|
|
mon = 0 if mon == -1 else mon
|
|
try:
|
|
monitor = monitors[mon]
|
|
except IndexError as exc:
|
|
msg = f"Monitor {mon!r} does not exist."
|
|
raise ScreenShotError(msg) from exc
|
|
|
|
output = output.format(mon=mon, date=datetime.now(UTC) if "{date" in output else None, **monitor)
|
|
if callable(callback):
|
|
callback(output)
|
|
sct = self.grab(monitor)
|
|
to_png(sct.rgb, sct.size, level=self.compression_level, output=output)
|
|
yield output
|
|
|
|
def shot(self, /, **kwargs: Any) -> str:
|
|
"""Helper to save the screenshot of the 1st monitor, by default.
|
|
You can pass the same arguments as for :meth:`save`.
|
|
"""
|
|
kwargs["mon"] = kwargs.get("mon", 1)
|
|
return next(self.save(**kwargs))
|
|
|
|
@staticmethod
|
|
def _merge(screenshot: ScreenShot, cursor: ScreenShot, /) -> ScreenShot:
|
|
"""Create composite image by blending screenshot and mouse cursor.
|
|
|
|
The cursor image should be in straight (not premultiplied) alpha.
|
|
"""
|
|
(cx, cy), (cw, ch) = cursor.pos, cursor.size
|
|
(x, y), (w, h) = screenshot.pos, screenshot.size
|
|
|
|
cx2, cy2 = cx + cw, cy + ch
|
|
x2, y2 = x + w, y + h
|
|
|
|
overlap = cx < x2 and cx2 > x and cy < y2 and cy2 > y
|
|
if not overlap:
|
|
return screenshot
|
|
|
|
screen_raw = screenshot.raw
|
|
cursor_raw = cursor.raw
|
|
|
|
cy, cy2 = (cy - y) * 4, (cy2 - y2) * 4
|
|
cx, cx2 = (cx - x) * 4, (cx2 - x2) * 4
|
|
start_count_y = -cy if cy < 0 else 0
|
|
start_count_x = -cx if cx < 0 else 0
|
|
stop_count_y = ch * 4 - max(cy2, 0)
|
|
stop_count_x = cw * 4 - max(cx2, 0)
|
|
rgb = range(3)
|
|
|
|
for count_y in range(start_count_y, stop_count_y, 4):
|
|
pos_s = (count_y + cy) * w + cx
|
|
pos_c = count_y * cw
|
|
|
|
for count_x in range(start_count_x, stop_count_x, 4):
|
|
spos = pos_s + count_x
|
|
cpos = pos_c + count_x
|
|
alpha = cursor_raw[cpos + 3]
|
|
|
|
if not alpha:
|
|
continue
|
|
|
|
if alpha == OPAQUE:
|
|
screen_raw[spos : spos + 3] = cursor_raw[cpos : cpos + 3]
|
|
else:
|
|
alpha2 = alpha / 255
|
|
for i in rgb:
|
|
screen_raw[spos + i] = int(cursor_raw[cpos + i] * alpha2 + screen_raw[spos + i] * (1 - alpha2))
|
|
|
|
return screenshot
|
|
|
|
@staticmethod
|
|
def _cfactory(
|
|
attr: Any,
|
|
func: str,
|
|
argtypes: list[Any],
|
|
restype: Any,
|
|
/,
|
|
errcheck: Callable | None = None,
|
|
) -> None:
|
|
"""Factory to create a ctypes function and automatically manage errors."""
|
|
meth = getattr(attr, func)
|
|
meth.argtypes = argtypes
|
|
meth.restype = restype
|
|
if errcheck:
|
|
meth.errcheck = errcheck
|
|
|
|
# Some backends may expose additional read-only attributes. Those
|
|
# are implemented here, as properties. By making them properties,
|
|
# instead of using __getattr__, they're also accessible to Sphinx
|
|
# and type checkers.
|
|
#
|
|
# Important: We need to be judicious in what we add here. We
|
|
# really don't want these to proliferate. Some, like
|
|
# max_displays, should probably be removed in 11.0. with_cursor
|
|
# should probably be moved to MSS instead of MSSImplementation (as
|
|
# noted there).
|
|
|
|
@property
|
|
def performance_status(self) -> list[str]:
|
|
"""Implementation-specific notes that might affect performance.
|
|
|
|
For instance, on GNU/Linux, when using the default XShmGetImage
|
|
backend, this will include a note if the MIT-SHM extension is
|
|
not usable.
|
|
|
|
This may not be ready until one screenshot has been taken.
|
|
|
|
This is meant only for debugging purposes; the contents are
|
|
subject to change at any time.
|
|
|
|
.. versionadded:: 10.2.0
|
|
"""
|
|
return self._impl.performance_status
|
|
|
|
@property
|
|
def max_displays(self) -> int:
|
|
"""Maximum number of displays to handle.
|
|
|
|
Availability: macOS
|
|
|
|
.. versionadded:: 8.0.0
|
|
"""
|
|
return self._impl.max_displays # type: ignore[attr-defined]
|
|
|
|
@property
|
|
def with_cursor(self) -> bool:
|
|
"""Include the mouse cursor in screenshots.
|
|
|
|
In some circumstances, it may not be possible to include the
|
|
cursor. In that case, MSS will automatically change this to
|
|
False when the object is created.
|
|
|
|
This cannot be changed after creating the object.
|
|
|
|
.. versionadded:: 8.0.0
|
|
"""
|
|
return self._impl.with_cursor
|
|
|
|
|
|
# TODO(jholveck): #493 Remove compatibility alias after 10.x transition period.
|
|
MSSBase = MSS
|