Skip to content

Commit

Permalink
Handle carriage return to support progress bar in logs
Browse files Browse the repository at this point in the history
This change breaks the read by reading each character and
yielding the breaks in <CR> or new lines instead of just waiting on new
line.

Studio PR: iterative/studio#10509
  • Loading branch information
amritghimire committed Aug 20, 2024
1 parent 70fe5a1 commit 4ba16a8
Showing 1 changed file with 15 additions and 3 deletions.
18 changes: 15 additions & 3 deletions src/datachain/catalog/catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,13 +120,25 @@ def noop(_: str):

@contextmanager
def print_and_capture(
stream: "IO[str]", callback: Callable[[str], None] = noop
stream: "IO[bytes]|IO[str]", callback: Callable[[str], None] = noop
) -> "Iterator[list[str]]":
lines: list[str] = []
append = lines.append

def loop() -> None:
for line in iter(stream.readline, ""):
buffer = b""
while byt := stream.read(1): # Read one byte at a time
buffer += byt.encode("utf-8") if isinstance(byt, str) else byt

if byt in (b"\n", b"\r"): # Check for newline or carriage return
line = buffer.decode("utf-8")
print(line, end="")
callback(line)
append(line)
buffer = b"" # Clear buffer for next line

if buffer: # Handle any remaining data in the buffer
line = buffer.decode("utf-8")
print(line, end="")
callback(line)
append(line)
Expand Down Expand Up @@ -2128,7 +2140,7 @@ def run_query(
stdout=subprocess.PIPE if capture_output else None,
stderr=subprocess.STDOUT if capture_output else None,
bufsize=1,
text=True,
text=False,
**kwargs,
) as proc:
os.close(w)
Expand Down

0 comments on commit 4ba16a8

Please sign in to comment.