client/experimental_client_with_swig/output_grabber.py

   1 import os
   2 import sys
   3 import threading
   4 import time
   5
   6 # From https://stackoverflow.com/a/29834357
   7 class OutputGrabber(object):
   8     """
   9     Class used to grab standard output or another stream.
  10     """
  11     escape_char = "\b"
  12
  13     def __init__(self, stream=None, threaded=False):
  14         self.origstream = stream
  15         self.threaded = threaded
  16         if self.origstream is None:
  17             self.origstream = sys.stdout
  18         self.origstreamfd = self.origstream.fileno()
  19         self.capturedtext = ""
  20         # Create a pipe so the stream can be captured:
  21         self.pipe_out, self.pipe_in = os.pipe()
  22
  23     def __enter__(self):
  24         self.start()
  25         return self
  26
  27     def __exit__(self, type, value, traceback):
  28         self.stop()
  29
  30     def start(self):
  31         """
  32         Start capturing the stream data.
  33         """
  34         self.capturedtext = ""
  35         # Save a copy of the stream:
  36         self.streamfd = os.dup(self.origstreamfd)
  37         # Replace the original stream with our write pipe:
  38         os.dup2(self.pipe_in, self.origstreamfd)
  39         if self.threaded:
  40             # Start thread that will read the stream:
  41             self.workerThread = threading.Thread(target=self.readOutput)
  42             self.workerThread.start()
  43             # Make sure that the thread is running and os.read() has executed:
  44             time.sleep(0.01)
  45
  46     def stop(self):
  47         """
  48         Stop capturing the stream data and save the text in `capturedtext`.
  49         """
  50         # Print the escape character to make the readOutput method stop:
  51         self.origstream.write(self.escape_char)
  52         # Flush the stream to make sure all our data goes in before
  53         # the escape character:
  54         self.origstream.flush()
  55         if self.threaded:
  56             # wait until the thread finishes so we are sure that
  57             # we have until the last character:
  58             self.workerThread.join()
  59         else:
  60             self.readOutput()
  61         # Close the pipe:
  62         os.close(self.pipe_in)
  63         os.close(self.pipe_out)
  64         # Restore the original stream:
  65         os.dup2(self.streamfd, self.origstreamfd)
  66         # Close the duplicate stream:
  67         os.close(self.streamfd)
  68
  69     def readOutput(self):
  70         """
  71         Read the stream data (one byte at a time)
  72         and save the text in `capturedtext`.
  73         """
  74         while True:
  75             char = os.read(self.pipe_out,1).decode(self.origstream.encoding)
  76             if not char or self.escape_char in char:
  77                 break
  78             self.capturedtext += char