__init__.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789
  1. # PyScreeze - PyScreeze is a simple, cross-platform screenshot module for Python 2 and 3.
  2. # By Al Sweigart al@inventwithpython.com
  3. __version__ = '1.0.1'
  4. import collections
  5. import datetime
  6. import functools
  7. import os
  8. import subprocess
  9. import sys
  10. import time
  11. import errno
  12. from contextlib import contextmanager
  13. from PIL import Image
  14. from PIL import ImageOps
  15. from PIL import ImageDraw
  16. from PIL import __version__ as PIL__version__
  17. from PIL import ImageGrab
  18. PILLOW_VERSION = tuple([int(x) for x in PIL__version__.split('.')])
  19. _useOpenCV: bool = False
  20. try:
  21. import cv2
  22. import numpy
  23. _useOpenCV = True
  24. except ImportError:
  25. pass # This is fine, useOpenCV will stay as False.
  26. RUNNING_PYTHON_2 = sys.version_info[0] == 2
  27. _PYGETWINDOW_UNAVAILABLE = True
  28. if sys.platform == 'win32':
  29. # On Windows, the monitor scaling can be set to something besides normal 100%.
  30. # PyScreeze and Pillow needs to account for this to make accurate screenshots.
  31. # TODO - How does macOS and Linux handle monitor scaling?
  32. import ctypes
  33. try:
  34. ctypes.windll.user32.SetProcessDPIAware()
  35. except AttributeError:
  36. pass # Windows XP doesn't support monitor scaling, so just do nothing.
  37. try:
  38. import pygetwindow
  39. except ImportError:
  40. _PYGETWINDOW_UNAVAILABLE = True
  41. else:
  42. _PYGETWINDOW_UNAVAILABLE = False
  43. GRAYSCALE_DEFAULT = True
  44. # For version 0.1.19 I changed it so that ImageNotFoundException was raised
  45. # instead of returning None. In hindsight, this change came too late, so I'm
  46. # changing it back to returning None. But I'm also including this option for
  47. # folks who would rather have it raise an exception.
  48. # For version 1.0.0, USE_IMAGE_NOT_FOUND_EXCEPTION is set to True by default.
  49. USE_IMAGE_NOT_FOUND_EXCEPTION = True
  50. GNOMESCREENSHOT_EXISTS = False
  51. try:
  52. if sys.platform.startswith('linux'):
  53. whichProc = subprocess.Popen(['which', 'gnome-screenshot'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  54. GNOMESCREENSHOT_EXISTS = whichProc.wait() == 0
  55. except OSError as ex:
  56. if ex.errno == errno.ENOENT:
  57. # if there is no "which" program to find gnome-screenshot, then assume there
  58. # is no gnome-screenshot.
  59. pass
  60. else:
  61. raise
  62. SCROT_EXISTS = False
  63. try:
  64. if sys.platform.startswith('linux'):
  65. whichProc = subprocess.Popen(['which', 'scrot'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
  66. SCROT_EXISTS = whichProc.wait() == 0
  67. except OSError as ex:
  68. if ex.errno == errno.ENOENT:
  69. # if there is no "which" program to find scrot, then assume there
  70. # is no scrot.
  71. pass
  72. else:
  73. raise
  74. # On Linux, figure out which window system is being used.
  75. if sys.platform.startswith('linux'):
  76. RUNNING_X11 = False
  77. RUNNING_WAYLAND = False
  78. if os.environ.get('XDG_SESSION_TYPE') == 'x11':
  79. RUNNING_X11 = True
  80. RUNNING_WAYLAND = False
  81. elif os.environ.get('XDG_SESSION_TYPE') == 'wayland':
  82. RUNNING_WAYLAND = True
  83. RUNNING_X11 = False
  84. elif 'WAYLAND_DISPLAY' in os.environ:
  85. RUNNING_WAYLAND = True
  86. RUNNING_X11 = False
  87. if sys.platform == 'win32':
  88. from ctypes import windll
  89. # win32 DC(DeviceContext) Manager
  90. @contextmanager
  91. def __win32_openDC(hWnd=0):
  92. """
  93. A context manager for handling calling GetDC() and ReleaseDC().
  94. This is used for win32 API calls, used by the pixel() function
  95. on Windows.
  96. Args:
  97. hWnd (int): The handle for the window to get a device context
  98. of, defaults to 0
  99. """
  100. hDC = windll.user32.GetDC(hWnd)
  101. if hDC == 0: # NULL
  102. raise WindowsError("windll.user32.GetDC failed : return NULL")
  103. try:
  104. yield hDC
  105. finally:
  106. windll.user32.ReleaseDC.argtypes = [ctypes.c_ssize_t, ctypes.c_ssize_t]
  107. if windll.user32.ReleaseDC(hWnd, hDC) == 0:
  108. raise WindowsError("windll.user32.ReleaseDC failed : return 0")
  109. Box = collections.namedtuple('Box', 'left top width height')
  110. Point = collections.namedtuple('Point', 'x y')
  111. RGB = collections.namedtuple('RGB', 'red green blue')
  112. class PyScreezeException(Exception):
  113. """PyScreezeException is a generic exception class raised when a
  114. PyScreeze-related error happens. If a PyScreeze function raises an
  115. exception that isn't PyScreezeException or a subclass, assume it is
  116. a bug in PyScreeze."""
  117. pass
  118. class ImageNotFoundException(PyScreezeException):
  119. """ImageNotFoundException is an exception class raised when the
  120. locate functions fail to locate an image. You must set
  121. pyscreeze.USE_IMAGE_NOT_FOUND_EXCEPTION to True to enable this feature.
  122. Otherwise, the locate functions will return None."""
  123. pass
  124. def requiresPyGetWindow(wrappedFunction):
  125. """
  126. A decorator that marks a function as requiring PyGetWindow to be installed.
  127. This raises PyScreezeException if Pillow wasn't imported.
  128. """
  129. @functools.wraps(wrappedFunction)
  130. def wrapper(*args, **kwargs):
  131. if _PYGETWINDOW_UNAVAILABLE:
  132. raise PyScreezeException('The PyGetWindow package is required to use this function.')
  133. return wrappedFunction(*args, **kwargs)
  134. return wrapper
  135. def _load_cv2(img, grayscale=None):
  136. """
  137. TODO
  138. """
  139. # load images if given filename, or convert as needed to opencv
  140. # Alpha layer just causes failures at this point, so flatten to RGB.
  141. # RGBA: load with -1 * cv2.CV_LOAD_IMAGE_COLOR to preserve alpha
  142. # to matchTemplate, need template and image to be the same wrt having alpha
  143. if grayscale is None:
  144. grayscale = GRAYSCALE_DEFAULT
  145. if isinstance(img, str):
  146. # The function imread loads an image from the specified file and
  147. # returns it. If the image cannot be read (because of missing
  148. # file, improper permissions, unsupported or invalid format),
  149. # the function returns an empty matrix
  150. # http://docs.opencv.org/3.0-beta/modules/imgcodecs/doc/reading_and_writing_images.html
  151. if grayscale:
  152. img_cv = cv2.imread(img, cv2.IMREAD_GRAYSCALE)
  153. else:
  154. img_cv = cv2.imread(img, cv2.IMREAD_COLOR)
  155. if img_cv is None:
  156. raise IOError(
  157. "Failed to read %s because file is missing, "
  158. "has improper permissions, or is an "
  159. "unsupported or invalid format" % img
  160. )
  161. elif isinstance(img, numpy.ndarray):
  162. # don't try to convert an already-gray image to gray
  163. if grayscale and len(img.shape) == 3: # and img.shape[2] == 3:
  164. img_cv = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
  165. else:
  166. img_cv = img
  167. elif hasattr(img, 'convert'):
  168. # assume its a PIL.Image, convert to cv format
  169. img_array = numpy.array(img.convert('RGB'))
  170. img_cv = img_array[:, :, ::-1].copy() # -1 does RGB -> BGR
  171. if grayscale:
  172. img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
  173. else:
  174. raise TypeError('expected an image filename, OpenCV numpy array, or PIL image')
  175. return img_cv
  176. def _locateAll_opencv(needleImage, haystackImage, grayscale=None, limit=10000, region=None, step=1, confidence=0.999):
  177. """
  178. TODO - rewrite this
  179. faster but more memory-intensive than pure python
  180. step 2 skips every other row and column = ~3x faster but prone to miss;
  181. to compensate, the algorithm automatically reduces the confidence
  182. threshold by 5% (which helps but will not avoid all misses).
  183. limitations:
  184. - OpenCV 3.x & python 3.x not tested
  185. - RGBA images are treated as RBG (ignores alpha channel)
  186. """
  187. if grayscale is None:
  188. grayscale = GRAYSCALE_DEFAULT
  189. confidence = float(confidence)
  190. needleImage = _load_cv2(needleImage, grayscale)
  191. needleHeight, needleWidth = needleImage.shape[:2]
  192. haystackImage = _load_cv2(haystackImage, grayscale)
  193. if region:
  194. haystackImage = haystackImage[region[1] : region[1] + region[3], region[0] : region[0] + region[2]]
  195. else:
  196. region = (0, 0) # full image; these values used in the yield statement
  197. if haystackImage.shape[0] < needleImage.shape[0] or haystackImage.shape[1] < needleImage.shape[1]:
  198. # avoid semi-cryptic OpenCV error below if bad size
  199. raise ValueError('needle dimension(s) exceed the haystack image or region dimensions')
  200. if step == 2:
  201. confidence *= 0.95
  202. needleImage = needleImage[::step, ::step]
  203. haystackImage = haystackImage[::step, ::step]
  204. else:
  205. step = 1
  206. # get all matches at once, credit: https://stackoverflow.com/questions/7670112/finding-a-subimage-inside-a-numpy-image/9253805#9253805
  207. result = cv2.matchTemplate(haystackImage, needleImage, cv2.TM_CCOEFF_NORMED)
  208. match_indices = numpy.arange(result.size)[(result > confidence).flatten()]
  209. matches = numpy.unravel_index(match_indices[:limit], result.shape)
  210. if len(matches[0]) == 0:
  211. if USE_IMAGE_NOT_FOUND_EXCEPTION:
  212. raise ImageNotFoundException('Could not locate the image (highest confidence = %.3f)' % result.max())
  213. else:
  214. return
  215. # use a generator for API consistency:
  216. matchx = matches[1] * step + region[0] # vectorized
  217. matchy = matches[0] * step + region[1]
  218. for x, y in zip(matchx, matchy):
  219. yield Box(x, y, needleWidth, needleHeight)
  220. def _locateAll_pillow(needleImage, haystackImage, grayscale=None, limit=None, region=None, step=1, confidence=None):
  221. """
  222. TODO
  223. """
  224. if confidence is not None:
  225. raise NotImplementedError('The confidence keyword argument is only available if OpenCV is installed.')
  226. # setup all the arguments
  227. if grayscale is None:
  228. grayscale = GRAYSCALE_DEFAULT
  229. needleFileObj = None
  230. if isinstance(needleImage, str):
  231. # 'image' is a filename, load the Image object
  232. needleFileObj = open(needleImage, 'rb')
  233. needleImage = Image.open(needleFileObj)
  234. haystackFileObj = None
  235. if isinstance(haystackImage, str):
  236. # 'image' is a filename, load the Image object
  237. haystackFileObj = open(haystackImage, 'rb')
  238. haystackImage = Image.open(haystackFileObj)
  239. if region is not None:
  240. haystackImage = haystackImage.crop((region[0], region[1], region[0] + region[2], region[1] + region[3]))
  241. else:
  242. region = (0, 0) # set to 0 because the code always accounts for a region
  243. if grayscale: # if grayscale mode is on, convert the needle and haystack images to grayscale
  244. needleImage = ImageOps.grayscale(needleImage)
  245. haystackImage = ImageOps.grayscale(haystackImage)
  246. else:
  247. # if not using grayscale, make sure we are comparing RGB images, not RGBA images.
  248. if needleImage.mode == 'RGBA':
  249. needleImage = needleImage.convert('RGB')
  250. if haystackImage.mode == 'RGBA':
  251. haystackImage = haystackImage.convert('RGB')
  252. # setup some constants we'll be using in this function
  253. needleWidth, needleHeight = needleImage.size
  254. haystackWidth, haystackHeight = haystackImage.size
  255. needleImageData = tuple(needleImage.getdata())
  256. haystackImageData = tuple(haystackImage.getdata())
  257. needleImageRows = [
  258. needleImageData[y * needleWidth : (y + 1) * needleWidth] for y in range(needleHeight)
  259. ] # LEFT OFF - check this
  260. needleImageFirstRow = needleImageRows[0]
  261. assert (
  262. len(needleImageFirstRow) == needleWidth
  263. ), 'The calculated width of first row of the needle image is not the same as the width of the image.'
  264. assert [len(row) for row in needleImageRows] == [
  265. needleWidth
  266. ] * needleHeight, 'The needleImageRows aren\'t the same size as the original image.'
  267. numMatchesFound = 0
  268. # NOTE: After running tests/benchmarks.py on the following code, it seem that having a step
  269. # value greater than 1 does not give *any* significant performance improvements.
  270. # Since using a step higher than 1 makes for less accurate matches, it will be
  271. # set to 1.
  272. step = 1 # hard-code step as 1 until a way to improve it can be figured out.
  273. if step == 1:
  274. firstFindFunc = _kmp
  275. else:
  276. firstFindFunc = _steppingFind
  277. for y in range(haystackHeight): # start at the leftmost column
  278. for matchx in firstFindFunc(
  279. needleImageFirstRow, haystackImageData[y * haystackWidth : (y + 1) * haystackWidth], step
  280. ):
  281. foundMatch = True
  282. for searchy in range(1, needleHeight, step):
  283. haystackStart = (searchy + y) * haystackWidth + matchx
  284. if (
  285. needleImageData[searchy * needleWidth : (searchy + 1) * needleWidth]
  286. != haystackImageData[haystackStart : haystackStart + needleWidth]
  287. ):
  288. foundMatch = False
  289. break
  290. if foundMatch:
  291. # Match found, report the x, y, width, height of where the matching region is in haystack.
  292. numMatchesFound += 1
  293. yield Box(matchx + region[0], y + region[1], needleWidth, needleHeight)
  294. if limit is not None and numMatchesFound >= limit:
  295. # Limit has been reached. Close file handles.
  296. if needleFileObj is not None:
  297. needleFileObj.close()
  298. if haystackFileObj is not None:
  299. haystackFileObj.close()
  300. return
  301. # There was no limit or the limit wasn't reached, but close the file handles anyway.
  302. if needleFileObj is not None:
  303. needleFileObj.close()
  304. if haystackFileObj is not None:
  305. haystackFileObj.close()
  306. if numMatchesFound == 0:
  307. if USE_IMAGE_NOT_FOUND_EXCEPTION:
  308. raise ImageNotFoundException('Could not locate the image.')
  309. else:
  310. return
  311. def locate(needleImage, haystackImage, **kwargs):
  312. """
  313. TODO
  314. """
  315. # Note: The gymnastics in this function is because we want to make sure to exhaust the iterator so that
  316. # the needle and haystack files are closed in locateAll.
  317. kwargs['limit'] = 1
  318. points = tuple(locateAll(needleImage, haystackImage, **kwargs))
  319. if len(points) > 0:
  320. return points[0]
  321. else:
  322. if USE_IMAGE_NOT_FOUND_EXCEPTION:
  323. raise ImageNotFoundException('Could not locate the image.')
  324. else:
  325. return None
  326. def locateOnScreen(image, minSearchTime=0, **kwargs):
  327. """TODO - rewrite this
  328. minSearchTime - amount of time in seconds to repeat taking
  329. screenshots and trying to locate a match. The default of 0 performs
  330. a single search.
  331. """
  332. start = time.time()
  333. while True:
  334. try:
  335. # the locateAll() function must handle cropping to return accurate coordinates,
  336. # so don't pass a region here.
  337. screenshotIm = screenshot(region=None)
  338. retVal = locate(image, screenshotIm, **kwargs)
  339. try:
  340. screenshotIm.fp.close()
  341. except AttributeError:
  342. # Screenshots on Windows won't have an fp since they came from
  343. # ImageGrab, not a file. Screenshots on Linux will have fp set
  344. # to None since the file has been unlinked
  345. pass
  346. if retVal or time.time() - start > minSearchTime:
  347. return retVal
  348. except ImageNotFoundException:
  349. if time.time() - start > minSearchTime:
  350. if USE_IMAGE_NOT_FOUND_EXCEPTION:
  351. raise
  352. else:
  353. return None
  354. def locateAllOnScreen(image, **kwargs):
  355. """
  356. TODO
  357. """
  358. # TODO - Should this raise an exception if zero instances of the image can be found on the
  359. # screen, instead of always returning a generator?
  360. # the locateAll() function must handle cropping to return accurate coordinates, so don't pass a region here.
  361. screenshotIm = screenshot(region=None)
  362. retVal = locateAll(image, screenshotIm, **kwargs)
  363. try:
  364. screenshotIm.fp.close()
  365. except AttributeError:
  366. # Screenshots on Windows won't have an fp since they came from
  367. # ImageGrab, not a file. Screenshots on Linux will have fp set
  368. # to None since the file has been unlinked
  369. pass
  370. return retVal
  371. def locateCenterOnScreen(image, **kwargs):
  372. """
  373. TODO
  374. """
  375. coords = locateOnScreen(image, **kwargs)
  376. if coords is None:
  377. return None
  378. else:
  379. return center(coords)
  380. def locateOnScreenNear(image, x, y):
  381. """
  382. TODO
  383. """
  384. foundMatchesBoxes = list(locateAllOnScreen(image))
  385. distancesSquared = [] # images[i] is related to distancesSquared[i]
  386. shortestDistanceIndex = 0 # The index of the shortest distance in `distances`
  387. # getting distance of all points from given point
  388. for foundMatchesBox in foundMatchesBoxes:
  389. foundMatchX, foundMatchY = center(foundMatchesBox)
  390. xDistance = abs(x - foundMatchX)
  391. yDistance = abs(y - foundMatchY)
  392. distancesSquared.append(xDistance * xDistance + yDistance * yDistance)
  393. if distancesSquared[-1] < distancesSquared[shortestDistanceIndex]:
  394. shortestDistanceIndex = len(distancesSquared) - 1
  395. # Returns the Box object of the match closest to x, y
  396. return foundMatchesBoxes[shortestDistanceIndex]
  397. def locateCenterOnScreenNear(image, x, y, **kwargs):
  398. """
  399. TODO
  400. """
  401. coords = locateOnScreenNear(image, x, y, **kwargs)
  402. if coords is None:
  403. return None
  404. else:
  405. return center(coords)
  406. @requiresPyGetWindow
  407. def locateOnWindow(image, title, **kwargs):
  408. """
  409. TODO
  410. """
  411. matchingWindows = pygetwindow.getWindowsWithTitle(title)
  412. if len(matchingWindows) == 0:
  413. raise PyScreezeException('Could not find a window with %s in the title' % (title))
  414. elif len(matchingWindows) > 1:
  415. raise PyScreezeException(
  416. 'Found multiple windows with %s in the title: %s' % (title, [str(win) for win in matchingWindows])
  417. )
  418. win = matchingWindows[0]
  419. win.activate()
  420. return locateOnScreen(image, region=(win.left, win.top, win.width, win.height), **kwargs)
  421. @requiresPyGetWindow
  422. def screenshotWindow(title):
  423. """
  424. TODO
  425. """
  426. pass # Not implemented yet.
  427. def showRegionOnScreen(region, outlineColor='red', filename='_showRegionOnScreen.png'):
  428. """
  429. TODO
  430. """
  431. # TODO - This function is useful! Document it!
  432. screenshotIm = screenshot()
  433. draw = ImageDraw.Draw(screenshotIm)
  434. region = (
  435. region[0],
  436. region[1],
  437. region[2] + region[0],
  438. region[3] + region[1],
  439. ) # convert from (left, top, right, bottom) to (left, top, width, height)
  440. draw.rectangle(region, outline=outlineColor)
  441. screenshotIm.save(filename)
  442. def _screenshot_win32(imageFilename=None, region=None, allScreens=False):
  443. """
  444. TODO
  445. """
  446. # TODO - Use the winapi to get a screenshot, and compare performance with ImageGrab.grab()
  447. # https://stackoverflow.com/a/3586280/1893164
  448. im = ImageGrab.grab(all_screens=allScreens)
  449. if region is not None:
  450. assert len(region) == 4, 'region argument must be a tuple of four ints'
  451. assert isinstance(region[0], int) and isinstance(region[1], int) and isinstance(region[2], int) and isinstance(region[3], int), 'region argument must be a tuple of four ints'
  452. im = im.crop((region[0], region[1], region[2] + region[0], region[3] + region[1]))
  453. if imageFilename is not None:
  454. im.save(imageFilename)
  455. return im
  456. def _screenshot_osx(imageFilename=None, region=None):
  457. """
  458. TODO
  459. """
  460. # TODO - use tmp name for this file.
  461. if PILLOW_VERSION < (6, 2, 1):
  462. # Use the screencapture program if Pillow is older than 6.2.1, which
  463. # is when Pillow supported ImageGrab.grab() on macOS. (It may have
  464. # supported it earlier than 6.2.1, but I haven't tested it.)
  465. if imageFilename is None:
  466. tmpFilename = 'screenshot%s.png' % (datetime.datetime.now().strftime('%Y-%m%d_%H-%M-%S-%f'))
  467. else:
  468. tmpFilename = imageFilename
  469. subprocess.call(['screencapture', '-x', tmpFilename])
  470. im = Image.open(tmpFilename)
  471. if region is not None:
  472. assert len(region) == 4, 'region argument must be a tuple of four ints'
  473. assert isinstance(region[0], int) and isinstance(region[1], int) and isinstance(region[2], int) and isinstance(region[3], int), 'region argument must be a tuple of four ints'
  474. im = im.crop((region[0], region[1], region[2] + region[0], region[3] + region[1]))
  475. os.unlink(tmpFilename) # delete image of entire screen to save cropped version
  476. im.save(tmpFilename)
  477. else:
  478. # force loading before unlinking, Image.open() is lazy
  479. im.load()
  480. if imageFilename is None:
  481. os.unlink(tmpFilename)
  482. else:
  483. # Use ImageGrab.grab() to get the screenshot if Pillow version 6.3.2 or later is installed.
  484. if region is not None:
  485. im = ImageGrab.grab(bbox=(region[0], region[1], region[2] + region[0], region[3] + region[1]))
  486. else:
  487. # Get full screen for screenshot
  488. im = ImageGrab.grab()
  489. return im
  490. def _screenshot_linux(imageFilename=None, region=None):
  491. """
  492. TODO
  493. """
  494. if imageFilename is None:
  495. tmpFilename = '.screenshot%s.png' % (datetime.datetime.now().strftime('%Y-%m%d_%H-%M-%S-%f'))
  496. else:
  497. tmpFilename = imageFilename
  498. # Version 9.2.0 introduced using gnome-screenshot for ImageGrab.grab()
  499. # on Linux, which is necessary to have screenshots work with Wayland
  500. # (the replacement for x11.) Therefore, for 3.7 and later, PyScreeze
  501. # uses/requires 9.2.0.
  502. if PILLOW_VERSION >= (9, 2, 0) and GNOMESCREENSHOT_EXISTS:
  503. # Pillow doesn't need tmpFilename because it works entirely in memory and doesn't
  504. # need to save an image file to disk.
  505. im = ImageGrab.grab() # use Pillow's grab() for Pillow 9.2.0 and later.
  506. if imageFilename is not None:
  507. im.save(imageFilename)
  508. if region is None:
  509. # Return the full screenshot.
  510. return im
  511. else:
  512. # Return just a region of the screenshot.
  513. assert len(region) == 4, 'region argument must be a tuple of four ints' # TODO fix this
  514. assert isinstance(region[0], int) and isinstance(region[1], int) and isinstance(region[2], int) and isinstance(region[3], int), 'region argument must be a tuple of four ints'
  515. im = im.crop((region[0], region[1], region[2] + region[0], region[3] + region[1]))
  516. return im
  517. elif RUNNING_X11 and SCROT_EXISTS: # scrot only runs on X11, not on Wayland.
  518. # Even if gnome-screenshot exists, use scrot on X11 because gnome-screenshot
  519. # has this annoying screen flash effect that you can't disable, but scrot does not.
  520. subprocess.call(['scrot', '-z', tmpFilename])
  521. elif GNOMESCREENSHOT_EXISTS: # gnome-screenshot runs on Wayland and X11.
  522. subprocess.call(['gnome-screenshot', '-f', tmpFilename])
  523. elif RUNNING_WAYLAND and SCROT_EXISTS and not GNOMESCREENSHOT_EXISTS:
  524. raise PyScreezeException(
  525. 'Your computer uses the Wayland window system. Scrot works on the X11 window system but not Wayland. You must install gnome-screenshot by running `sudo apt install gnome-screenshot`' # noqa
  526. )
  527. else:
  528. raise Exception(
  529. 'To take screenshots, you must install Pillow version 9.2.0 or greater and gnome-screenshot by running `sudo apt install gnome-screenshot`' # noqa
  530. )
  531. im = Image.open(tmpFilename)
  532. if region is not None:
  533. assert len(region) == 4, 'region argument must be a tuple of four ints'
  534. assert isinstance(region[0], int) and isinstance(region[1], int) and isinstance(region[2], int) and isinstance(region[3], int), 'region argument must be a tuple of four ints'
  535. im = im.crop((region[0], region[1], region[2] + region[0], region[3] + region[1]))
  536. os.unlink(tmpFilename) # delete image of entire screen to save cropped version
  537. im.save(tmpFilename)
  538. else:
  539. # force loading before unlinking, Image.open() is lazy
  540. im.load()
  541. if imageFilename is None:
  542. os.unlink(tmpFilename)
  543. return im
  544. def _kmp(needle, haystack, _dummy): # Knuth-Morris-Pratt search algorithm implementation (to be used by screen capture)
  545. """
  546. TODO
  547. """
  548. # build table of shift amounts
  549. shifts = [1] * (len(needle) + 1)
  550. shift = 1
  551. for pos in range(len(needle)):
  552. while shift <= pos and needle[pos] != needle[pos - shift]:
  553. shift += shifts[pos - shift]
  554. shifts[pos + 1] = shift
  555. # do the actual search
  556. startPos = 0
  557. matchLen = 0
  558. for c in haystack:
  559. while matchLen == len(needle) or matchLen >= 0 and needle[matchLen] != c:
  560. startPos += shifts[matchLen]
  561. matchLen -= shifts[matchLen]
  562. matchLen += 1
  563. if matchLen == len(needle):
  564. yield startPos
  565. def _steppingFind(needle, haystack, step):
  566. """
  567. TODO
  568. """
  569. for startPos in range(0, len(haystack) - len(needle) + 1):
  570. foundMatch = True
  571. for pos in range(0, len(needle), step):
  572. if haystack[startPos + pos] != needle[pos]:
  573. foundMatch = False
  574. break
  575. if foundMatch:
  576. yield startPos
  577. def center(coords):
  578. """
  579. Returns a `Point` object with the x and y set to an integer determined by the format of `coords`.
  580. The `coords` argument is a 4-integer tuple of (left, top, width, height).
  581. For example:
  582. >>> center((10, 10, 6, 8))
  583. Point(x=13, y=14)
  584. >>> center((10, 10, 7, 9))
  585. Point(x=13, y=14)
  586. >>> center((10, 10, 8, 10))
  587. Point(x=14, y=15)
  588. """
  589. # TODO - one day, add code to handle a Box namedtuple.
  590. return Point(coords[0] + int(coords[2] / 2), coords[1] + int(coords[3] / 2))
  591. def pixelMatchesColor(x, y, expectedRGBColor, tolerance=0):
  592. """
  593. Return True if the pixel at x, y is matches the expected color of the RGB
  594. tuple, each color represented from 0 to 255, within an optional tolerance.
  595. """
  596. # TODO DEPRECATE THIS FUNCTION
  597. # Note: Automate the Boring Stuff 2nd edition documented that you could call
  598. # pixelMatchesColor((x, y), rgb) instead of pixelMatchesColor(x, y, rgb).
  599. # Lets correct that for the 1.0 release.
  600. if isinstance(x, collections.abc.Sequence) and len(x) == 2:
  601. raise TypeError('pixelMatchesColor() has updated and no longer accepts a tuple of (x, y) values for the first argument. Pass these arguments as two separate arguments instead: pixelMatchesColor(x, y, rgb) instead of pixelMatchesColor((x, y), rgb)')
  602. pix = pixel(x, y)
  603. if len(pix) == 3 or len(expectedRGBColor) == 3: # RGB mode
  604. r, g, b = pix[:3]
  605. exR, exG, exB = expectedRGBColor[:3]
  606. return (abs(r - exR) <= tolerance) and (abs(g - exG) <= tolerance) and (abs(b - exB) <= tolerance)
  607. elif len(pix) == 4 and len(expectedRGBColor) == 4: # RGBA mode
  608. r, g, b, a = pix
  609. exR, exG, exB, exA = expectedRGBColor
  610. return (
  611. (abs(r - exR) <= tolerance)
  612. and (abs(g - exG) <= tolerance)
  613. and (abs(b - exB) <= tolerance)
  614. and (abs(a - exA) <= tolerance)
  615. )
  616. else:
  617. assert False, (
  618. 'Color mode was expected to be length 3 (RGB) or 4 (RGBA), but pixel is length %s and expectedRGBColor is length %s' # noqa
  619. % (len(pix), len(expectedRGBColor))
  620. )
  621. def pixel(x, y):
  622. """
  623. Returns the color of the screen pixel at x, y as an RGB tuple, each color represented from 0 to 255.
  624. """
  625. # Note: Automate the Boring Stuff 2nd edition documented that you could call
  626. # pixel((x, y), rgb) instead of pixel(x, y, rgb).
  627. # Lets correct that for the 1.0 release.
  628. if isinstance(x, collections.abc.Sequence) and len(x) == 2:
  629. raise TypeError('pixel() has updated and no longer accepts a tuple of (x, y) values for the first argument. Pass these arguments as two separate arguments instead: pixel(x, y) instead of pixel((x, y))')
  630. if sys.platform == 'win32':
  631. # On Windows, calling GetDC() and GetPixel() is twice as fast as using our screenshot() function.
  632. with __win32_openDC() as hdc: # handle will be released automatically
  633. color = windll.gdi32.GetPixel(hdc, x, y)
  634. if color < 0:
  635. raise WindowsError("windll.gdi32.GetPixel failed : return {}".format(color))
  636. # color is in the format 0xbbggrr https://msdn.microsoft.com/en-us/library/windows/desktop/dd183449(v=vs.85).aspx
  637. bbggrr = "{:0>6x}".format(color) # bbggrr => 'bbggrr' (hex)
  638. b, g, r = (int(bbggrr[i : i + 2], 16) for i in range(0, 6, 2))
  639. return (r, g, b)
  640. else:
  641. # Need to select only the first three values of the color in
  642. # case the returned pixel has an alpha channel
  643. return RGB(*(screenshot().getpixel((x, y))[:3]))
  644. # set the screenshot() function based on the platform running this module
  645. if sys.platform == 'darwin':
  646. screenshot = _screenshot_osx
  647. elif sys.platform == 'win32':
  648. screenshot = _screenshot_win32
  649. elif sys.platform.startswith('linux'):
  650. # Everything else is considered to be Linux.
  651. screenshot = _screenshot_linux
  652. else:
  653. raise NotImplementedError('PyScreeze is not supported on platform ' + sys.platform)
  654. # set the locateAll function to use opencv if possible; python 3 needs opencv 3.0+
  655. # TODO - Should this raise an exception if zero instances of the image can be found
  656. # on the screen, instead of always returning a generator?
  657. locateAll = _locateAll_pillow
  658. if _useOpenCV:
  659. locateAll = _locateAll_opencv
  660. if not RUNNING_PYTHON_2 and cv2.__version__ < '3':
  661. locateAll = _locateAll_pillow