add line buffering to pyio

Author: Christopher Doris

docs/src/releasenotes.md

@@ -1,5 +1,9 @@
 # Release Notes
 
+## Unreleased
+* Adds `line_buffering` option to `PyIO`.
+* Improvements to stdout when using `juliacall.ipython` includng line-buffering.
+
 ## 0.9.7 (2022-10-11)
 * If CondaPkg is using the Null backend, PythonCall will now use `python` from the PATH.
 * Bug fixes.

pysrc/juliacall/ipython.py

@@ -34,8 +34,8 @@ def load_ipython_extension(ip):
     PythonCall.seval("""begin
         const _redirected_stdout = redirect_stdout()
         const _redirected_stderr = redirect_stderr()
-        const _py_stdout = PyIO(pyimport("sys" => "stdout"))
-        const _py_stderr = PyIO(pyimport("sys" => "stderr"))
+        const _py_stdout = PyIO(pyimport("sys" => "stdout"); line_buffering=true)
+        const _py_stderr = PyIO(pyimport("sys" => "stderr"); line_buffering=true)
         const _redirect_stdout_task = @async write($_py_stdout, $_redirected_stdout)
         const _redirect_stderr_task = @async write($_py_stderr, $_redirected_stderr)
         function _flush_stdio()

src/pywrap/PyIO.jl

@@ -1,5 +1,5 @@
 """
-    PyIO(x; own=false, text=missing, buflen=4096)
+    PyIO(x; own=false, text=missing, line_buffering=false, buflen=4096)
 
 Wrap the Python IO stream `x` as a Julia IO stream.
 
@@ -10,6 +10,8 @@ If `text=true` then `x` must be a text stream and only UTF-8 must be written (i.
 If `text` is not specified then it is chosen automatically.
 If `x` is a text stream and you really need a binary stream, then often `PyIO(x.buffer)` will work.
 
+If `line_buffering=true` then output is flushed at each line.
+
 For efficiency, reads and writes are buffered before being sent to `x`.
 The size of the buffers is `buflen`.
 The buffers are cleared using `flush`.
@@ -20,6 +22,8 @@ mutable struct PyIO <: IO
     own::Bool
     # true if `o` is text, false if binary
     text::Bool
+    # true to flush whenever '\n' or '\r' is encountered
+    line_buffering::Bool
     # true if we are definitely at the end of the file; false if we are not or don't know
     eof::Bool
     # input buffer
@@ -29,7 +33,7 @@ mutable struct PyIO <: IO
     obuflen::Int
     obuf::Vector{UInt8}
 
-    function PyIO(x; own::Bool = false, text::Union{Missing,Bool} = missing, buflen::Integer = 4096, ibuflen::Integer = buflen, obuflen::Integer = buflen)
+    function PyIO(x; own::Bool = false, text::Union{Missing,Bool} = missing, buflen::Integer = 4096, ibuflen::Integer = buflen, obuflen::Integer = buflen, line_buffering::Bool=false)
         if text === missing
             text = pyhasattr(x, "encoding")
         end
@@ -39,7 +43,7 @@ mutable struct PyIO <: IO
         ibuflen > 0 || error("ibuflen must be positive")
         obuflen = convert(Int, obuflen)
         obuflen > 0 || error("obuflen must be positive")
-        new(Py(x), own, text, false, ibuflen, UInt8[], obuflen, UInt8[])
+        new(Py(x), own, text, line_buffering, false, ibuflen, UInt8[], obuflen, UInt8[])
     end
 end
 export PyIO
@@ -136,23 +140,37 @@ Base.isopen(io::PyIO) = !pyconvert(Bool, @py io.closed)
 
 function Base.unsafe_write(io::PyIO, ptr::Ptr{UInt8}, n::UInt)
     ntodo = n
-    while true
+    while ntodo > 0
         nroom = max(0, io.obuflen - length(io.obuf))
         if ntodo < nroom
-            append!(io.obuf, unsafe_wrap(Array, ptr, ntodo))
-            return n
+            buf = unsafe_wrap(Array, ptr, ntodo)
+            if io.line_buffering
+                i = findlast(∈((0x0A, 0x0D)), buf)
+                if i === nothing
+                    append!(io.obuf, buf)
+                else
+                    append!(io.obuf, unsafe_wrap(Array, ptr, i))
+                    putobuf(io)
+                    append!(io.obuf, unsafe_wrap(Array, ptr+i, ntodo-i))
+                end
+            else
+                append!(io.obuf, buf)
+            end
+            break
         else
-            append!(io.obuf, unsafe_wrap(Array, ptr, nroom))
+            buf = unsafe_wrap(Array, ptr, nroom)
+            append!(io.obuf, buf)
             putobuf(io)
             ptr += nroom
             ntodo -= nroom
         end
     end
+    return n
 end
 
 function Base.write(io::PyIO, c::UInt8)
     push!(io.obuf, c)
-    if length(io.obuf) ≥ io.obuflen
+    if (length(io.obuf) ≥ io.obuflen) || (io.line_buffering && (c == 0x0A || c == 0x0D))
         putobuf(io)
     end
     return