WIP: Enable ImageGrab grab for Linux

Tests/test_imagegrab.py

@@ -2,67 +2,47 @@
 import sys
 
 import pytest
+from PIL import ImageGrab
 
 from .helper import assert_image
 
-try:
- from PIL import ImageGrab
 
- class TestImageGrab:
- def test_grab(self):
- for im in [
- ImageGrab.grab(),
- ImageGrab.grab(include_layered_windows=True),
- ImageGrab.grab(all_screens=True),
+class TestImageGrab:
+ def test_grab(self):
+ native_support = sys.platform in ["darwin", "win32"]
+ if native_support or ImageGrab._has_imagemagick():
+ for args in [
+ {},
+ {"include_layered_windows": True},
+ {"all_screens": True},
+ {"bbox": (10, 20, 50, 80)},
  ]:
- assert_image(im, im.mode, im.size)
-
- im = ImageGrab.grab(bbox=(10, 20, 50, 80))
- assert_image(im, im.mode, (40, 60))
-
- def test_grabclipboard(self):
- if sys.platform == "darwin":
- subprocess.call(["screencapture", "-cx"])
- else:
- p = subprocess.Popen(
- ["powershell", "-command", "-"], stdin=subprocess.PIPE
- )
- p.stdin.write(
- b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
+ try:
+ im = ImageGrab.grab(**args)
+ except IOError as e:
+ if not native_support and str(e) == "Unable to open X server":
+ continue
+ else:
+ raise
+ assert_image(im, im.mode, (40, 60) if "bbox" in args else im.size)
+ else:
+ pytest.raises(IOError, ImageGrab.grab)
+
+ def test_grabclipboard(self):
+ if sys.platform == "darwin":
+ subprocess.call(["screencapture", "-cx"])
+ elif sys.platform == "win32":
+ p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
+ p.stdin.write(
+ b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
 [Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
 $bmp = New-Object Drawing.Bitmap 200, 200
 [Windows.Forms.Clipboard]::SetImage($bmp)"""
- )
- p.communicate()
-
- im = ImageGrab.grabclipboard()
- assert_image(im, im.mode, im.size)
-
-
-except ImportError:
-
- class TestImageGrab:
- @pytest.mark.skip(reason="ImageGrab ImportError")
- def test_skip(self):
- pass
-
-
-class TestImageGrabImport:
- def test_import(self):
- # Arrange
- exception = None
-
- # Act
- try:
- from PIL import ImageGrab
-
- ImageGrab.__name__ # dummy to prevent Pyflakes warning
- except Exception as e:
- exception = e
-
- # Assert
- if sys.platform in ["win32", "darwin"]:
- assert exception is None
+ )
+ p.communicate()
  else:
- assert isinstance(exception, ImportError)
- assert str(exception) == "ImageGrab is macOS and Windows only"
+ pytest.raises(NotImplementedError, ImageGrab.grabclipboard)
+ return
+
+ im = ImageGrab.grabclipboard()
+ assert_image(im, im.mode, im.size)

src/PIL/ImageGrab.py

@@ -22,8 +22,16 @@
 
 from . import Image
 
-if sys.platform not in ["win32", "darwin"]:
- raise ImportError("ImageGrab is macOS and Windows only")
+
+def _has_imagemagick():
+ try:
+ with open(os.devnull, "wb") as devnull:
+ subprocess.check_call(["import", "--version"], stdout=devnull)
+ return True
+ except OSError:
+ # No ImageMagick
+ pass
+ return False
 
 
 def grab(bbox=None, include_layered_windows=False, all_screens=False):
@@ -38,8 +46,9 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False):
  im_cropped = im.crop(bbox)
  im.close()
  return im_cropped
- else:
- offset, size, data = Image.core.grabscreen(include_layered_windows, all_screens)
+ elif sys.platform == "win32":
+ grabber = Image.core.grabscreen
+ offset, size, data = grabber(include_layered_windows, all_screens)
  im = Image.frombytes(
  "RGB",
  size,
@@ -54,6 +63,22 @@ def grab(bbox=None, include_layered_windows=False, all_screens=False):
  x0, y0 = offset
  left, top, right, bottom = bbox
  im = im.crop((left - x0, top - y0, right - x0, bottom - y0))
+ else:
+ if not _has_imagemagick:
+ raise IOError("grab requires ImageMagick unless used on macOS or Windows")
+ fh, filepath = tempfile.mkstemp(".png")
+ os.close(fh)
+
+ p = subprocess.Popen(
+ ["import", "-window", "root", filepath], stderr=subprocess.PIPE
+ )
+ if b"unable to open X server" in p.communicate()[1]:
+ raise IOError("Unable to open X server")
+ im = Image.open(filepath)
+ im.load()
+ os.unlink(filepath)
+ if bbox:
+ im = im.crop(bbox)
  return im
 
 
@@ -81,11 +106,15 @@ def grabclipboard():
  im.load()
  os.unlink(filepath)
  return im
- else:
+
+ elif sys.platform == "win32":
  data = Image.core.grabclipboard()
  if isinstance(data, bytes):
  from . import BmpImagePlugin
  import io
 
  return BmpImagePlugin.DibImageFile(io.BytesIO(data))
  return data
+
+ else:
+ raise NotImplementedError("grabclipboard is macOS and Windows only")