forked from BoboTiG/python-mss
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbase.py
More file actions
196 lines (155 loc) · 6.27 KB
/
base.py
File metadata and controls
196 lines (155 loc) · 6.27 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
"""
This is part of the MSS Python's module.
Source: https://github.com/BoboTiG/python-mss
"""
from abc import ABCMeta, abstractmethod
from datetime import datetime
from typing import TYPE_CHECKING
from threading import Lock
from .exception import ScreenShotError
from .screenshot import ScreenShot
from .tools import to_png
if TYPE_CHECKING:
# pylint: disable=ungrouped-imports
from typing import Any, Callable, Iterator, List, Optional, Type # noqa
from .models import Monitor, Monitors # noqa
lock = Lock()
class MSSBase(metaclass=ABCMeta):
""" This class will be overloaded by a system specific one. """
__slots__ = {"_monitors", "cls_image", "compression_level"}
def __init__(self):
self.cls_image = ScreenShot # type: Type[ScreenShot]
self.compression_level = 6
self._monitors = [] # type: Monitors
def __enter__(self):
# type: () -> MSSBase
""" For the cool call `with MSS() as mss:`. """
return self
def __exit__(self, *_):
""" For the cool call `with MSS() as mss:`. """
self.close()
@abstractmethod
def _grab_impl(self, monitor):
# type: (Monitor) -> ScreenShot
"""
Retrieve all pixels from a monitor. Pixels have to be RGB.
That method has to be run using a threading lock.
"""
@abstractmethod
def _monitors_impl(self):
# type: () -> None
"""
Get positions of monitors (has to be run using a threading lock).
It must populate self._monitors.
"""
def close(self):
# type: () -> None
""" Clean-up. """
def grab(self, monitor):
# type: (Monitor) -> ScreenShot
"""
Retrieve screen pixels for a given monitor.
Note: *monitor* can be a tuple like PIL.Image.grab() accepts.
:param monitor: The coordinates and size of the box to capture.
See :meth:`monitors <monitors>` for object details.
:return :class:`ScreenShot <ScreenShot>`.
"""
# 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],
}
with lock:
return self._grab_impl(monitor)
@property
def monitors(self):
# type: () -> 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
}
"""
if not self._monitors:
with lock:
self._monitors_impl()
return self._monitors
def save(self, mon=0, output="monitor-{mon}.png", callback=None):
# type: (int, str, Callable[[str], None]) -> Iterator[str]
"""
Grab a screen shot and save it to a file.
:param int mon: The monitor to screen shot (default=0).
-1: grab one screen shot of all monitors
0: grab one screen shot by monitor
N: grab the screen shot of the monitor N
:param str output: The output filename.
It can take several keywords to customize the filename:
- `{mon}`: the monitor number
- `{top}`: the screen shot y-coordinate of the upper-left corner
- `{left}`: the screen shot x-coordinate of the upper-left corner
- `{width}`: the screen shot's width
- `{height}`: the screen shot's height
- `{date}`: the current date using the default formatter
As it is using the `format()` function, you can specify
formatting options like `{date:%Y-%m-%s}`.
:param callable callback: Callback called before saving the
screen shot to a file. Take the `output` argument as parameter.
:return generator: Created file(s).
"""
monitors = self.monitors
if not monitors:
raise ScreenShotError("No monitor found.")
if mon == 0:
# One screen shot by monitor
for idx, monitor in enumerate(monitors[1:], 1):
fname = output.format(mon=idx, date=datetime.now(), **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 screen shot of all monitors together or
# a screen shot of the monitor N.
mon = 0 if mon == -1 else mon
try:
monitor = monitors[mon]
except IndexError:
# pylint: disable=raise-missing-from
raise ScreenShotError("Monitor {!r} does not exist.".format(mon))
output = output.format(mon=mon, date=datetime.now(), **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):
# type: (Any) -> str
"""
Helper to save the screen shot of the 1st monitor, by default.
You can pass the same arguments as for ``save``.
"""
kwargs["mon"] = kwargs.get("mon", 1)
return next(self.save(**kwargs))
@staticmethod
def _cfactory(attr, func, argtypes, restype, errcheck=None):
# type: (Any, str, List[Any], Any, Optional[Callable]) -> 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