88import sys
99import threading
1010import unicodedata
11- from typing import Any , BinaryIO , Iterable , List , Optional , TextIO , Union
11+ from typing import Any , Iterable , List , Optional , TextIO , Union
1212
1313from 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
378400class 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:
500519class 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