Skip to content

Commit

Permalink
Handle carriage return to support progress bar in logs (#326)
Browse files Browse the repository at this point in the history
  • Loading branch information
amritghimire authored Aug 22, 2024
1 parent 20b42d6 commit dfa2068
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 dfa2068

Please sign in to comment.