Pillow

1
from __future__ import annotations
2

3
import os
4
import shutil
5
import subprocess
6
import sys
7

8
import pytest
9

10
from PIL import Image, ImageGrab
11

12
from .helper import assert_image_equal_tofile, skip_unless_feature
13

14

15
class TestImageGrab:
16
    @pytest.mark.skipif(
17
        os.environ.get("USERNAME") == "ContainerAdministrator",
18
        reason="can't grab screen when running in Docker",
19
    )
20
    @pytest.mark.skipif(
21
        sys.platform not in ("win32", "darwin"), reason="requires Windows or macOS"
22
    )
23
    def test_grab(self) -> None:
24
        ImageGrab.grab()
25
        ImageGrab.grab(include_layered_windows=True)
26
        ImageGrab.grab(all_screens=True)
27

28
        im = ImageGrab.grab(bbox=(10, 20, 50, 80))
29
        assert im.size == (40, 60)
30

31
    @skip_unless_feature("xcb")
32
    def test_grab_x11(self) -> None:
33
        try:
34
            if sys.platform not in ("win32", "darwin"):
35
                ImageGrab.grab()
36

37
            ImageGrab.grab(xdisplay="")
38
        except OSError as e:
39
            pytest.skip(str(e))
40

41
    @pytest.mark.skipif(Image.core.HAVE_XCB, reason="tests missing XCB")
42
    def test_grab_no_xcb(self) -> None:
43
        if sys.platform not in ("win32", "darwin") and not shutil.which(
44
            "gnome-screenshot"
45
        ):
46
            with pytest.raises(OSError) as e:
47
                ImageGrab.grab()
48
            assert str(e.value).startswith("Pillow was built without XCB support")
49

50
        with pytest.raises(OSError) as e:
51
            ImageGrab.grab(xdisplay="")
52
        assert str(e.value).startswith("Pillow was built without XCB support")
53

54
    @skip_unless_feature("xcb")
55
    def test_grab_invalid_xdisplay(self) -> None:
56
        with pytest.raises(OSError) as e:
57
            ImageGrab.grab(xdisplay="error.test:0.0")
58
        assert str(e.value).startswith("X connection failed")
59

60
    def test_grabclipboard(self) -> None:
61
        if sys.platform == "darwin":
62
            subprocess.call(["screencapture", "-cx"])
63

64
            ImageGrab.grabclipboard()
65
        elif sys.platform == "win32":
66
            p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
67
            p.stdin.write(
68
                b"""[Reflection.Assembly]::LoadWithPartialName("System.Drawing")
69
[Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
70
$bmp = New-Object Drawing.Bitmap 200, 200
71
[Windows.Forms.Clipboard]::SetImage($bmp)"""
72
            )
73
            p.communicate()
74

75
            ImageGrab.grabclipboard()
76
        else:
77
            if not shutil.which("wl-paste") and not shutil.which("xclip"):
78
                with pytest.raises(
79
                    NotImplementedError,
80
                    match="wl-paste or xclip is required for"
81
                    r" ImageGrab.grabclipboard\(\) on Linux",
82
                ):
83
                    ImageGrab.grabclipboard()
84

85
    @pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
86
    def test_grabclipboard_file(self) -> None:
87
        p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
88
        assert p.stdin is not None
89
        p.stdin.write(rb'Set-Clipboard -Path "Tests\images\hopper.gif"')
90
        p.communicate()
91

92
        im = ImageGrab.grabclipboard()
93
        assert isinstance(im, list)
94
        assert len(im) == 1
95
        assert os.path.samefile(im[0], "Tests/images/hopper.gif")
96

97
    @pytest.mark.skipif(sys.platform != "win32", reason="Windows only")
98
    def test_grabclipboard_png(self) -> None:
99
        p = subprocess.Popen(["powershell", "-command", "-"], stdin=subprocess.PIPE)
100
        assert p.stdin is not None
101
        p.stdin.write(
102
            rb"""$bytes = [System.IO.File]::ReadAllBytes("Tests\images\hopper.png")
103
$ms = new-object System.IO.MemoryStream(, $bytes)
104
[Reflection.Assembly]::LoadWithPartialName("System.Windows.Forms")
105
[Windows.Forms.Clipboard]::SetData("PNG", $ms)"""
106
        )
107
        p.communicate()
108

109
        im = ImageGrab.grabclipboard()
110
        assert isinstance(im, Image.Image)
111
        assert_image_equal_tofile(im, "Tests/images/hopper.png")
112

113
    @pytest.mark.skipif(
114
        (
115
            sys.platform != "linux"
116
            or not all(shutil.which(cmd) for cmd in ("wl-paste", "wl-copy"))
117
        ),
118
        reason="Linux with wl-clipboard only",
119
    )
120
    @pytest.mark.parametrize("ext", ("gif", "png", "ico"))
121
    def test_grabclipboard_wl_clipboard(self, ext: str) -> None:
122
        image_path = "Tests/images/hopper." + ext
123
        with open(image_path, "rb") as fp:
124
            subprocess.call(["wl-copy"], stdin=fp)
125
        im = ImageGrab.grabclipboard()
126
        assert isinstance(im, Image.Image)
127
        assert_image_equal_tofile(im, image_path)
128

129
    @pytest.mark.skipif(
130
        (
131
            sys.platform != "linux"
132
            or not all(shutil.which(cmd) for cmd in ("wl-paste", "wl-copy"))
133
        ),
134
        reason="Linux with wl-clipboard only",
135
    )
136
    @pytest.mark.parametrize("arg", ("text", "--clear"))
137
    def test_grabclipboard_wl_clipboard_errors(self, arg: str) -> None:
138
        subprocess.call(["wl-copy", arg])
139
        assert ImageGrab.grabclipboard() is None
140