diff --git a/README.md b/README.md index abfc4f72..0a133f20 100644 --- a/README.md +++ b/README.md @@ -16,18 +16,17 @@ with mss() as sct: An ultra-fast cross-platform multiple screenshots module in pure python using ctypes. - **Python 3.9+**, PEP8 compliant, no dependency, thread-safe; -- very basic, it will grab one screenshot by monitor or a screenshot of all monitors and save it to a PNG file; +- very basic, it will grab one screenshot by monitor or window, or a screenshot of all monitors and save it to a PNG file; - but you can use PIL and benefit from all its formats (or add yours directly); - integrate well with Numpy and OpenCV; - it could be easily embedded into games and other software which require fast and platform optimized methods to grab screenshots (like AI, Computer Vision); - get the [source code on GitHub](https://github.com/BoboTiG/python-mss); - learn with a [bunch of examples](https://python-mss.readthedocs.io/examples.html); - you can [report a bug](https://github.com/BoboTiG/python-mss/issues); -- need some help? Use the tag *python-mss* on [Stack Overflow](https://stackoverflow.com/questions/tagged/python-mss); +- need some help? Use the tag _python-mss_ on [Stack Overflow](https://stackoverflow.com/questions/tagged/python-mss); - and there is a [complete, and beautiful, documentation](https://python-mss.readthedocs.io) :) - **MSS** stands for Multiple ScreenShots; - ## Installation You can install it with pip: diff --git a/src/mss/__main__.py b/src/mss/__main__.py index 384ad344..6da510fd 100644 --- a/src/mss/__main__.py +++ b/src/mss/__main__.py @@ -31,7 +31,9 @@ def main(*args: str) -> int: help="the PNG compression level", ) cli_args.add_argument("-m", "--monitor", default=0, type=int, help="the monitor to screenshot") - cli_args.add_argument("-o", "--output", default="monitor-{mon}.png", help="the output file name") + cli_args.add_argument("-w", "--window", default=None, help="the window to screenshot") + cli_args.add_argument("-p", "--process", default=None, help="the process to screenshot") + cli_args.add_argument("-o", "--output", default=None, help="the output file name") cli_args.add_argument("--with-cursor", default=False, action="store_true", help="include the cursor") cli_args.add_argument( "-q", @@ -43,7 +45,8 @@ def main(*args: str) -> int: cli_args.add_argument("-v", "--version", action="version", version=__version__) options = cli_args.parse_args(args or None) - kwargs = {"mon": options.monitor, "output": options.output} + output = options.output or ("window-{win}.png" if options.window or options.process else "monitor-{mon}.png") + kwargs = {"mon": options.monitor, "output": output, "win": options.window, "proc": options.process} if options.coordinates: try: top, left, width, height = options.coordinates.split(",") diff --git a/src/mss/base.py b/src/mss/base.py index 8a7397f5..12efe45f 100644 --- a/src/mss/base.py +++ b/src/mss/base.py @@ -16,7 +16,7 @@ if TYPE_CHECKING: # pragma: nocover from collections.abc import Callable, Iterator - from mss.models import Monitor, Monitors + from mss.models import Monitor, Monitors, Window, Windows try: from datetime import UTC @@ -34,7 +34,7 @@ class MSSBase(metaclass=ABCMeta): """This class will be overloaded by a system specific one.""" - __slots__ = {"_monitors", "cls_image", "compression_level", "with_cursor"} + __slots__ = {"_monitors", "_windows", "cls_image", "compression_level", "with_cursor"} def __init__( self, @@ -51,6 +51,7 @@ def __init__( self.compression_level = compression_level self.with_cursor = with_cursor self._monitors: Monitors = [] + self._windows: Windows = [] def __enter__(self) -> MSSBase: # noqa:PYI034 """For the cool call `with MSS() as mss:`.""" @@ -70,12 +71,24 @@ def _grab_impl(self, monitor: Monitor, /) -> ScreenShot: That method has to be run using a threading lock. """ + @abstractmethod + def _grab_window_impl(self, window: Window, /) -> ScreenShot: + """Retrieve all pixels from a window. Pixels have to be RGB. + That method has to be run using a threading lock. + """ + @abstractmethod def _monitors_impl(self) -> None: """Get positions of monitors (has to be run using a threading lock). It must populate self._monitors. """ + @abstractmethod + def _windows_impl(self) -> None: + """Get ids of windows (has to be run using a threading lock). + It must populate self._windows. + """ + def close(self) -> None: # noqa:B027 """Clean-up.""" @@ -103,6 +116,31 @@ def grab(self, monitor: Monitor | tuple[int, int, int, int], /) -> ScreenShot: return self._merge(screenshot, cursor) return screenshot + def grab_window( + self, window: Window | str | None = None, /, *, name: str | None = None, process: str | None = None + ) -> ScreenShot: + """Retrieve screen pixels for a given window. + + :param window: The window to capture or its name. + See :meth:`windows ` for object details. + :param str name: The window name. + :param str process: The window process name. + :return :class:`ScreenShot `. + """ + if isinstance(window, str): + name = window + window = None + + if window is None: + windows = self.find_windows(name, process) + if not windows: + msg = f"Window {window!r} not found." + raise ScreenShotError(msg) + window = windows[0] + + with lock: + return self._grab_window_impl(window) + @property def monitors(self) -> Monitors: """Get positions of all monitors. @@ -128,11 +166,54 @@ def monitors(self) -> Monitors: return self._monitors + @property + def windows(self) -> Windows: + """Get ids, names, and proceesses of all windows. + Unlike monitors, this method does not use a cache, as the list of + windows can change at any time. + + Each window is a dict with: + { + 'id': the window id or handle, + 'name': the window name, + 'process': the window process name, + 'bounds': the window bounds as 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 + } + } + """ + with lock: + self._windows_impl() + + return self._windows + + def find_windows(self, name: str | None = None, process: str | None = None) -> Windows: + """Find windows by name and/or process name. + + :param str name: The window name. + :param str process: The window process name. + :return list: List of windows. + """ + windows = self.windows + if name is None and process is None: + return windows + if name is None: + return [window for window in windows if window["process"] == process] + if process is None: + return [window for window in windows if window["name"] == name] + return [window for window in windows if window["name"] == name and window["process"] == process] + def save( self, /, *, mon: int = 0, + win: str | None = None, + proc: str | None = None, output: str = "monitor-{mon}.png", callback: Callable[[str], None] | None = None, ) -> Iterator[str]: @@ -166,7 +247,21 @@ def save( msg = "No monitor found." raise ScreenShotError(msg) - if mon == 0: + if win or proc: + windows = self.find_windows(win, proc) + if not windows: + msg = f"Window {(win or proc)!r} not found." + raise ScreenShotError(msg) + window = windows[0] + + fname = output.format(win=win or proc, date=datetime.now(UTC) if "{date" in output else None) + if callable(callback): + callback(fname) + + sct = self.grab_window(window) + to_png(sct.rgb, sct.size, level=self.compression_level, output=fname) + yield fname + elif 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) diff --git a/src/mss/darwin.py b/src/mss/darwin.py index a56e05a8..0d7ac016 100644 --- a/src/mss/darwin.py +++ b/src/mss/darwin.py @@ -7,7 +7,20 @@ import ctypes import ctypes.util import sys -from ctypes import POINTER, Structure, c_double, c_float, c_int32, c_ubyte, c_uint32, c_uint64, c_void_p +from ctypes import ( + POINTER, + Structure, + c_bool, + c_char_p, + c_double, + c_float, + c_int32, + c_long, + c_ubyte, + c_uint32, + c_uint64, + c_void_p, +) from platform import mac_ver from typing import TYPE_CHECKING, Any @@ -16,7 +29,7 @@ from mss.screenshot import ScreenShot, Size if TYPE_CHECKING: # pragma: nocover - from mss.models import CFunctions, Monitor + from mss.models import CConstants, CFunctions, Monitor, Window __all__ = ("MSS",) @@ -78,6 +91,28 @@ def __repr__(self) -> str: "CGRectStandardize": ("core", [CGRect], CGRect), "CGRectUnion": ("core", [CGRect, CGRect], CGRect), "CGWindowListCreateImage": ("core", [CGRect, c_uint32, c_uint32, c_uint32], c_void_p), + "CGWindowListCopyWindowInfo": ("core", [c_uint32, c_uint32], c_void_p), + "CFArrayGetCount": ("core", [c_void_p], c_uint64), + "CFArrayGetValueAtIndex": ("core", [c_void_p, c_uint64], c_void_p), + "CFNumberGetValue": ("core", [c_void_p, c_int32, c_void_p], c_bool), + "CFStringGetCString": ("core", [c_void_p, c_char_p, c_long, c_uint32], c_bool), + "CFDictionaryGetValue": ("core", [c_void_p, c_void_p], c_void_p), + "CGRectMakeWithDictionaryRepresentation": ("core", [c_void_p, POINTER(CGRect)], c_bool), +} + +CCONSTANTS: CConstants = { + # Syntax: cconstant: type or value + "kCGWindowNumber": c_void_p, + "kCGWindowName": c_void_p, + "kCGWindowOwnerName": c_void_p, + "kCGWindowBounds": c_void_p, + "kCGWindowListOptionOnScreenOnly": 0b0001, + "kCGWindowListOptionIncludingWindow": 0b1000, + "kCFStringEncodingUTF8": 0x08000100, + "kCGNullWindowID": 0, + "kCFNumberSInt32Type": 3, + "kCGWindowImageBoundsIgnoreFraming": 0b0001, + "CGRectNull": CGRect, } @@ -86,7 +121,7 @@ class MSS(MSSBase): It uses intensively the CoreGraphics library. """ - __slots__ = {"core", "max_displays"} + __slots__ = {"constants", "core", "max_displays"} def __init__(self, /, **kwargs: Any) -> None: """MacOS initialisations.""" @@ -96,6 +131,7 @@ def __init__(self, /, **kwargs: Any) -> None: self._init_library() self._set_cfunctions() + self._set_cconstants() def _init_library(self) -> None: """Load the CoreGraphics library.""" @@ -118,6 +154,15 @@ def _set_cfunctions(self) -> None: for func, (attr, argtypes, restype) in CFUNCTIONS.items(): cfactory(attrs[attr], func, argtypes, restype) + def _set_cconstants(self) -> None: + """Set all ctypes constants and attach them to attributes.""" + self.constants = {} + for name, value in CCONSTANTS.items(): + if isinstance(value, type) and hasattr(value, "in_dll"): + self.constants[name] = value.in_dll(self.core, name) + else: + self.constants[name] = value + def _monitors_impl(self) -> None: """Get positions of monitors. It will populate self._monitors.""" int_ = int @@ -165,16 +210,61 @@ def _monitors_impl(self) -> None: "height": int_(all_monitors.size.height), } - def _grab_impl(self, monitor: Monitor, /) -> ScreenShot: - """Retrieve all pixels from a monitor. Pixels have to be RGB.""" + def _windows_impl(self) -> None: core = self.core - rect = CGRect((monitor["left"], monitor["top"]), (monitor["width"], monitor["height"])) + constants = self.constants + kCGWindowListOptionOnScreenOnly = constants["kCGWindowListOptionOnScreenOnly"] # noqa: N806 + kCFNumberSInt32Type = constants["kCFNumberSInt32Type"] # noqa: N806 + kCGWindowNumber = constants["kCGWindowNumber"] # noqa: N806 + kCGWindowName = constants["kCGWindowName"] # noqa: N806 + kCGWindowOwnerName = constants["kCGWindowOwnerName"] # noqa: N806 + kCGWindowBounds = constants["kCGWindowBounds"] # noqa: N806 + kCFStringEncodingUTF8 = constants["kCFStringEncodingUTF8"] # noqa: N806 + + window_list = core.CGWindowListCopyWindowInfo(kCGWindowListOptionOnScreenOnly, 0) + + window_count = core.CFArrayGetCount(window_list) + + str_buf = ctypes.create_string_buffer(256) + self._windows = [] + for i in range(window_count): + window_info = core.CFArrayGetValueAtIndex(window_list, i) + window_id = c_int32() + core.CFNumberGetValue( + core.CFDictionaryGetValue(window_info, kCGWindowNumber), kCFNumberSInt32Type, ctypes.byref(window_id) + ) - image_ref = core.CGWindowListCreateImage(rect, 1, 0, 0) - if not image_ref: - msg = "CoreGraphics.CGWindowListCreateImage() failed." - raise ScreenShotError(msg) + core.CFStringGetCString( + core.CFDictionaryGetValue(window_info, kCGWindowName), str_buf, 256, kCFStringEncodingUTF8 + ) + window_name = str_buf.value.decode("utf-8") + core.CFStringGetCString( + core.CFDictionaryGetValue(window_info, kCGWindowOwnerName), str_buf, 256, kCFStringEncodingUTF8 + ) + process_name = str_buf.value.decode("utf-8") + + window_bound_ref = core.CFDictionaryGetValue(window_info, kCGWindowBounds) + window_bounds = CGRect() + core.CGRectMakeWithDictionaryRepresentation(window_bound_ref, ctypes.byref(window_bounds)) + + self._windows.append( + { + "id": window_id.value, + "name": window_name, + "process": process_name, + "bounds": { + "left": int(window_bounds.origin.x), + "top": int(window_bounds.origin.y), + "width": int(window_bounds.size.width), + "height": int(window_bounds.size.height), + }, + } + ) + + def _image_to_data(self, image_ref: c_void_p, /) -> bytearray: + """Convert a CGImageRef to a bytearray.""" + core = self.core width = core.CGImageGetWidth(image_ref) height = core.CGImageGetHeight(image_ref) prov = copy_data = None @@ -198,14 +288,53 @@ def _grab_impl(self, monitor: Monitor, /) -> ScreenShot: end = start + width * bytes_per_pixel cropped.extend(data[start:end]) data = cropped + + return data finally: if prov: core.CGDataProviderRelease(prov) if copy_data: core.CFRelease(copy_data) + def _grab_impl(self, monitor: Monitor, /) -> ScreenShot: + """Retrieve all pixels from a monitor. Pixels have to be RGB.""" + core = self.core + rect = CGRect((monitor["left"], monitor["top"]), (monitor["width"], monitor["height"])) + + image_ref = core.CGWindowListCreateImage(rect, 1, 0, 0) + if not image_ref: + msg = "CoreGraphics.CGWindowListCreateImage() failed." + raise ScreenShotError(msg) + + width = core.CGImageGetWidth(image_ref) + height = core.CGImageGetHeight(image_ref) + data = self._image_to_data(image_ref) + return self.cls_image(data, monitor, size=Size(width, height)) + def _grab_window_impl(self, window: Window, /) -> ScreenShot: + """Retrieve all pixels from a window. Pixels have to be RGB.""" + core = self.core + constants = self.constants + bounds = window["bounds"] + + rect = constants["CGRectNull"] + list_option = constants["kCGWindowListOptionIncludingWindow"] + window_id = window["id"] + image_option = constants["kCGWindowImageBoundsIgnoreFraming"] + + image_ref = core.CGWindowListCreateImage(rect, list_option, window_id, image_option) + + if not image_ref: + msg = "CoreGraphics.CGWindowListCreateImage() failed." + raise ScreenShotError(msg) + + width = core.CGImageGetWidth(image_ref) + height = core.CGImageGetHeight(image_ref) + data = self._image_to_data(image_ref) + + return self.cls_image(data, bounds, size=Size(width, height)) + def _cursor_impl(self) -> ScreenShot | None: """Retrieve all cursor data. Pixels have to be RGB.""" return None diff --git a/src/mss/models.py b/src/mss/models.py index 665a41bc..fea515af 100644 --- a/src/mss/models.py +++ b/src/mss/models.py @@ -7,10 +7,14 @@ Monitor = dict[str, int] Monitors = list[Monitor] +Window = dict[str, Any] +Windows = list[Window] + Pixel = tuple[int, int, int] Pixels = list[tuple[Pixel, ...]] CFunctions = dict[str, tuple[str, list[Any], Any]] +CConstants = dict[str, Any] class Pos(NamedTuple): diff --git a/src/tests/test_find_windows.py b/src/tests/test_find_windows.py new file mode 100644 index 00000000..7d8824ff --- /dev/null +++ b/src/tests/test_find_windows.py @@ -0,0 +1,34 @@ +"""This is part of the MSS Python's module. +Source: https://github.com/BoboTiG/python-mss. +""" + +from mss import mss + + +def test_get_windows() -> None: + with mss() as sct: + assert sct.windows + + +def test_find_windows_by_name() -> None: + with mss() as sct: + source_window = sct.windows[0] + target_window = sct.find_windows(name=source_window["name"])[0] + assert source_window["name"] == target_window["name"] + assert source_window["process"] == target_window["process"] + assert source_window["bounds"]["top"] == target_window["bounds"]["top"] + assert source_window["bounds"]["left"] == target_window["bounds"]["left"] + assert source_window["bounds"]["width"] == target_window["bounds"]["width"] + assert source_window["bounds"]["height"] == target_window["bounds"]["height"] + + +def test_find_windows_by_process() -> None: + with mss() as sct: + source_window = sct.windows[0] + target_window = sct.find_windows(process=source_window["process"])[0] + assert source_window["name"] == target_window["name"] + assert source_window["process"] == target_window["process"] + assert source_window["bounds"]["top"] == target_window["bounds"]["top"] + assert source_window["bounds"]["left"] == target_window["bounds"]["left"] + assert source_window["bounds"]["width"] == target_window["bounds"]["width"] + assert source_window["bounds"]["height"] == target_window["bounds"]["height"] diff --git a/src/tests/test_save.py b/src/tests/test_save.py index 9597206c..58c6ed7c 100644 --- a/src/tests/test_save.py +++ b/src/tests/test_save.py @@ -81,3 +81,17 @@ def test_output_format_date_custom() -> None: filename = sct.shot(mon=1, output=fmt) assert filename == fmt.format(date=datetime.now(tz=UTC)) assert Path(filename).is_file() + +def test_output_format_window_name() -> None: + with mss(display=os.getenv("DISPLAY")) as sct: + window = sct.windows[0] + filename = next(sct.save(win=window["name"], output="window-{win}.png")) + assert filename == f"window-{window["name"]}.png" + assert Path(filename).is_file() + +def test_output_format_window_process() -> None: + with mss(display=os.getenv("DISPLAY")) as sct: + window = sct.windows[0] + filename = next(sct.save(proc=window["process"], output="process-{win}.png")) + assert filename == f"process-{window["process"]}.png" + assert Path(filename).is_file() \ No newline at end of file