Skip to content

Commit 9117fdc

Browse files
committed
Removing support for wrapping binary streams since cmd2 assumes stdout will be a text stream in all of its code like poutput()
1 parent 7a20c2e commit 9117fdc

File tree

1 file changed

+29
-10
lines changed

1 file changed

+29
-10
lines changed

cmd2/utils.py

Lines changed: 29 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import sys
99
import threading
1010
import unicodedata
11-
from typing import Any, BinaryIO, Iterable, List, Optional, TextIO, Union
11+
from typing import Any, Iterable, List, Optional, TextIO, Union
1212

1313
from wcwidth import wcswidth
1414

@@ -297,7 +297,7 @@ def __init__(self, inner_stream, echo: bool = False,
297297
encoding: str = 'utf-8', errors: str = 'replace') -> None:
298298
"""
299299
Initializer
300-
:param inner_stream: the emulated stream
300+
:param inner_stream: the wrapped stream. Should be a TextIO or StdSim instance.
301301
:param echo: if True, then all input will be echoed to inner_stream
302302
:param encoding: codec for encoding/decoding strings (defaults to utf-8)
303303
:param errors: how to handle encoding/decoding errors (defaults to replace)
@@ -350,6 +350,17 @@ def isatty(self) -> bool:
350350
else:
351351
return False
352352

353+
@property
354+
def line_buffering(self) -> bool:
355+
"""
356+
Handle when the inner stream doesn't have a line_buffering attribute which is the case
357+
when running unit tests because pytest sets stdout to a pytest EncodedFile object.
358+
"""
359+
try:
360+
return self.inner_stream.line_buffering
361+
except AttributeError:
362+
return False
363+
353364
def __getattr__(self, item: str):
354365
if item in self.__dict__:
355366
return self.__dict__[item]
@@ -361,6 +372,9 @@ class ByteBuf(object):
361372
"""
362373
Used by StdSim to write binary data and stores the actual bytes written
363374
"""
375+
# Used to know when to flush the StdSim
376+
NEWLINES = [b'\n', b'\r']
377+
364378
def __init__(self, std_sim_instance: StdSim) -> None:
365379
self.byte_buf = bytearray()
366380
self.std_sim_instance = std_sim_instance
@@ -374,14 +388,22 @@ def write(self, b: bytes) -> None:
374388
if self.std_sim_instance.echo:
375389
self.std_sim_instance.inner_stream.buffer.write(b)
376390

391+
# Since StdSim wraps TextIO streams, we will flush the stream if line buffering is on
392+
# and the bytes being written contain a new line character. This is helpful when StdSim
393+
# is being used to capture output of a shell command because it causes the output to print
394+
# to the screen more often than if we waited for the stream to flush its buffer.
395+
if self.std_sim_instance.line_buffering:
396+
if any(newline in b for newline in ByteBuf.NEWLINES):
397+
self.std_sim_instance.flush()
398+
377399

378400
class ProcReader(object):
379401
"""
380402
Used to captured stdout and stderr from a Popen process if any of those were set to subprocess.PIPE.
381403
If neither are pipes, then the process will run normally and no output will be captured.
382404
"""
383-
def __init__(self, proc: subprocess.Popen, stdout: Union[StdSim, BinaryIO, TextIO],
384-
stderr: Union[StdSim, BinaryIO, TextIO]) -> None:
405+
def __init__(self, proc: subprocess.Popen, stdout: Union[StdSim, TextIO],
406+
stderr: Union[StdSim, TextIO]) -> None:
385407
"""
386408
ProcReader initializer
387409
:param proc: the Popen process being read from
@@ -457,17 +479,14 @@ def _reader_thread_func(self, read_stdout: bool) -> None:
457479
self._write_bytes(write_stream, available)
458480

459481
@staticmethod
460-
def _write_bytes(stream: Union[StdSim, BinaryIO, TextIO], to_write: bytes) -> None:
482+
def _write_bytes(stream: Union[StdSim, TextIO], to_write: bytes) -> None:
461483
"""
462484
Write bytes to a stream
463485
:param stream: the stream being written to
464486
:param to_write: the bytes being written
465487
"""
466488
try:
467-
if hasattr(stream, 'buffer'):
468-
stream.buffer.write(to_write)
469-
else:
470-
stream.write(to_write)
489+
stream.buffer.write(to_write)
471490
except BrokenPipeError:
472491
# This occurs if output is being piped to a process that closed
473492
pass
@@ -500,7 +519,7 @@ def __exit__(self, *args) -> None:
500519
class RedirectionSavedState(object):
501520
"""Created by each command to store information about their redirection."""
502521

503-
def __init__(self, self_stdout: Union[StdSim, BinaryIO, TextIO], sys_stdout: Union[StdSim, BinaryIO, TextIO],
522+
def __init__(self, self_stdout: Union[StdSim, TextIO], sys_stdout: Union[StdSim, TextIO],
504523
pipe_proc_reader: Optional[ProcReader]) -> None:
505524
# Used to restore values after the command ends
506525
self.saved_self_stdout = self_stdout

0 commit comments

Comments
 (0)