Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

accept fileno-having objects for GIL-less capture #83

Merged
merged 1 commit into from
Apr 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,22 @@ with pipes(logger, stderr=STDOUT):
call_some_c_function()
```

Forward C-level output to a file (avoids GIL issues with a background thread, new in 3.1):

```python
from wurlitzer import pipes, STDOUT

with open("log.txt", "ab") as f, pipes(f, stderr=STDOUT):
blocking_gil_holding_function()
```

Or even simpler, enable it as an IPython extension:

```
%load_ext wurlitzer
```

To forward all C-level output to IPython during execution.
To forward all C-level output to IPython (e.g. Jupyter cell output) during execution.

## Acknowledgments

Expand Down
36 changes: 36 additions & 0 deletions test.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,3 +206,39 @@ def test_log_pipes(caplog):
# check 'stream' extra
assert record.stream
assert record.name == "wurlitzer." + record.stream


def test_two_file_pipes(tmpdir):

test_stdout = tmpdir / "stdout.txt"
test_stderr = tmpdir / "stderr.txt"

with test_stdout.open("ab") as stdout_f, test_stderr.open("ab") as stderr_f:
w = Wurlitzer(stdout_f, stderr_f)
with w:
assert w.thread is None
printf("some stdout")
printf_err("some stderr")

with test_stdout.open() as f:
assert f.read() == "some stdout\n"
with test_stderr.open() as f:
assert f.read() == "some stderr\n"


def test_one_file_pipe(tmpdir):

test_stdout = tmpdir / "stdout.txt"

with test_stdout.open("ab") as stdout_f:
stderr = io.StringIO()
w = Wurlitzer(stdout_f, stderr)
with w as (stdout, stderr):
assert w.thread is not None
printf("some stdout")
printf_err("some stderr")
assert not w.thread.is_alive()

with test_stdout.open() as f:
assert f.read() == "some stdout\n"
assert stderr.getvalue() == "some stderr\n"
48 changes: 36 additions & 12 deletions wurlitzer.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,6 +210,17 @@ def _setup_pipe(self, name):
save_fd = os.dup(real_fd)
self._save_fds[name] = save_fd

try:
capture_fd = getattr(self, "_" + name).fileno()
except Exception:
pass
else:
# if it has a fileno(),
# dup directly to capture file,
# no pipes needed
dup2(capture_fd, real_fd)
return None

pipe_out, pipe_in = os.pipe()
# set max pipe buffer size (linux only)
if self._bufsize:
Expand Down Expand Up @@ -272,19 +283,32 @@ def __enter__(self):
self._flush()
# setup handle
self._setup_handle()
self._control_r, self._control_w = os.pipe()

# create pipe for stdout
pipes = [self._control_r]
names = {self._control_r: 'control'}
pipes = []
names = {}
if self._stdout:
pipe = self._setup_pipe('stdout')
pipes.append(pipe)
names[pipe] = 'stdout'
if pipe:
pipes.append(pipe)
names[pipe] = 'stdout'
if self._stderr:
pipe = self._setup_pipe('stderr')
pipes.append(pipe)
names[pipe] = 'stderr'
if pipe:
pipes.append(pipe)
names[pipe] = 'stderr'

if not pipes:
# no pipes to handle (e.g. direct FD capture)
# so no forwarder thread needed
self.thread = None
return self.handle

# setup forwarder thread

self._control_r, self._control_w = os.pipe()
pipes.append(self._control_r)
names[self._control_r] = "control"

# flush pipes in a background thread to avoid blocking
# the reader thread when the buffer is full
Expand Down Expand Up @@ -366,11 +390,11 @@ def forwarder():
def __exit__(self, exc_type, exc_value, traceback):
# flush before exiting
self._flush()

# signal output is complete on control pipe
os.write(self._control_w, b'\1')
self.thread.join()
os.close(self._control_w)
if self.thread:
# signal output is complete on control pipe
os.write(self._control_w, b'\1')
self.thread.join()
os.close(self._control_w)

# restore original state
for name, real_fd in self._real_fds.items():
Expand Down
Loading