Skip to content

Commit

Permalink
Merge pull request #83 from minrk/file_pipes
Browse files Browse the repository at this point in the history
accept fileno-having objects for GIL-less capture
  • Loading branch information
minrk authored Apr 24, 2024
2 parents 38caf37 + 21cd32b commit 628a8dc
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 13 deletions.
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

0 comments on commit 628a8dc

Please sign in to comment.