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
|
#!/usr/bin/env python3
# Corey Goldberg, 2012-2025
# License: MIT
"""Run a headless display inside X virtual framebuffer (Xvfb)"""
import os
import platform
import shutil
import subprocess
import tempfile
import time
try:
import fcntl
except ImportError:
system = platform.system()
raise OSError(f"xvfbwrapper is not supported on this platform: {system}")
from random import randint
class Xvfb:
# Maximum value to use for a display. 32-bit maxint is the
# highest Xvfb currently supports
MAX_DISPLAY = 2147483647
def __init__(
self,
width=800,
height=680,
colordepth=24,
tempdir=None,
display=None,
environ=None,
timeout=10,
**kwargs,
):
self.width = width
self.height = height
self.colordepth = colordepth
self._tempdir = tempdir or tempfile.gettempdir()
self._timeout = timeout
self.new_display = display
self.environ = environ if environ else os.environ
if not self._xvfb_exists():
raise OSError("Can't find Xvfb. Please install it and try again")
self.xvfb_cmd = []
self.extra_xvfb_args = [
"-screen",
"0",
f"{self.width}x{self.height}x{self.colordepth}",
]
for key, value in kwargs.items():
self.extra_xvfb_args += [f"-{key}", value]
if "DISPLAY" in self.environ:
self.orig_display_var = self.environ["DISPLAY"]
else:
self.orig_display_var = None
self.proc = None
def __enter__(self) -> "Xvfb":
self.start()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.stop()
def start(self) -> None:
if self.new_display is not None:
if not self._get_lock_for_display(self.new_display):
raise ValueError(f"Could not lock display :{self.new_display}")
else:
self.new_display = self._get_next_unused_display()
display_var = f":{self.new_display}"
self.xvfb_cmd = ["Xvfb", display_var] + self.extra_xvfb_args
self.proc = subprocess.Popen(
self.xvfb_cmd,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
close_fds=True,
)
start = time.time()
while not self._local_display_exists(self.new_display):
time.sleep(1e-3)
if time.time() - start > self._timeout:
self.stop()
raise RuntimeError(f"Xvfb display did not open: {self.xvfb_cmd}")
ret_code = self.proc.poll()
if ret_code is None:
self._set_display(display_var)
else:
self._cleanup_lock_file()
raise RuntimeError(f"Xvfb did not start ({ret_code}): {self.xvfb_cmd}")
def stop(self) -> None:
try:
if self.orig_display_var is None:
self.environ.pop("DISPLAY", None)
else:
self._set_display(self.orig_display_var)
if self.proc is not None:
try:
self.proc.terminate()
self.proc.wait(self._timeout)
except OSError:
pass
self.proc = None
finally:
self._cleanup_lock_file()
def _xvfb_exists(self) -> bool:
"""Check that Xvfb is available on PATH and is executable."""
return True if shutil.which("Xvfb") is not None else False
def _cleanup_lock_file(self):
"""
This should always get called if the process exits safely
with Xvfb.stop() (whether called explicitly, or by __exit__).
If you are ending up with /tmp/X123-lock files when Xvfb is not
running, then Xvfb is not exiting cleanly. Always either call
Xvfb.stop() in a finally block, or use Xvfb as a context manager
to ensure lock files are purged.
"""
self._lock_display_file.close()
try:
os.remove(self._lock_display_file.name)
except OSError:
pass
def _get_lock_for_display(self, display) -> bool:
"""
In order to ensure multi-process safety, this method attempts
to acquire an exclusive lock on a temporary file whose name
contains the display number for Xvfb.
"""
tempfile_path = os.path.join(self._tempdir, f".X{display}-lock")
try:
self._lock_display_file = open(tempfile_path, "w")
except PermissionError:
return False
else:
try:
fcntl.flock(self._lock_display_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
return False
else:
return True
def _get_next_unused_display(self) -> int:
"""
Randomly chooses a display number and tries to acquire a lock for this
number. If the lock could be acquired, returns this number, otherwise
choses a new one.
:return: free display number
"""
while True:
rand = randint(1, self.__class__.MAX_DISPLAY)
if self._get_lock_for_display(rand):
return rand
else:
continue
def _local_display_exists(self, display) -> bool:
return os.path.exists(f"/tmp/.X11-unix/X{display}")
def _set_display(self, display_var):
self.environ["DISPLAY"] = display_var
|
