Source Code for Module TEES.Utils.Stream

  1  """ 
  2  Capture stdout or stderr messages to a log file 
  3   
  4  Based on code by Greg Pinero (Primary Searcher) 
  5   
  6  Capture print statments and write them to a log file 
  7  but still allow them to be printed on the screen. 
  8  """ 
  9  import sys, os 
 10  import time 
 11  import types 
 12   
 13 -class StreamModifier: 
 14      """ 
 15      This class implements a write-method and can therefore replace a stream 
 16      such as sys.stderr or sys.stdout. The write method first writes the text 
 17      to a log file, then passes it on to the original stream. 
 18      """ 
 19 -    def __init__(self, stream): 
 20          self.stream = stream 
 21          self.logfiles = [] 
 22          self.logfilenames = [] 
 23          self.indent = None 
 24          self.timeStamp = None 
 25          self.timeStampDuplicates = False 
 26          self.prevTime = None 
 27          self.newLine = True 
 28          self.buffer = "" 
 29       
 30 -    def setLog(self, logfile=None): 
 31          if logfile != None: 
 32              self.logfiles = [logfile] 
 33              self.logfilenames = [logfile.name] 
 34          else: 
 35              self.logfiles = [] 
 36              self.logfilenames = [] 
 37       
 38 -    def addLog(self, logfile): 
 39          self.logfiles.append(logfile) 
 40          self.logfilenames.append(logfile.name) 
 41       
 42 -    def removeLog(self, logfileName, streamName): 
 43          logfilesToKeep = [] 
 44          logfilenamesToKeep = [] 
 45          removed = None 
 46          removedName = None 
 47          for i in range(len(self.logfiles)): 
 48              if self.logfilenames[i] != logfileName: 
 49                  logfilesToKeep.append(self.logfiles[i]) 
 50                  logfilenamesToKeep.append(self.logfilenames[i]) 
 51              else: 
 52                  removed = self.logfiles[i] 
 53                  removedName = self.logfilenames[i] 
 54          self.logfiles = logfilesToKeep 
 55          self.logfilenames = logfilenamesToKeep 
 56           
 57          if removedName != None: 
 58              print >> sys.stderr, "Stopped logging", streamName, "to", removedName 
 59          else: 
 60              print >> sys.stderr, "Log not open for ", streamName + ":", logfileName 
 61          return removed 
 62       
 63 -    def setIndent(self, indent=None): 
 64          self.indent = indent 
 65       
 66 -    def setTimeStamp(self, format=None, duplicates=False): 
 67          self.timeStamp = format 
 68          self.timeStampDuplicates = duplicates 
 69       
 70 -    def writeToStream(self, text): 
 71          """ 
 72          Write directly to the stream without adding to the log file 
 73          """ 
 74          self.stream.write(text) 
 75       
 76 -    def writeToLog(self, text, filename): 
 77          """ 
 78          Write directly to the log file without sending the input to the stream 
 79          """ 
 80          for logfile in self.logfiles: 
 81              if filename == None or logfile.name == filename: 
 82                  logfile.write(text) 
 83                  logfile.flush() 
 84       
 85 -    def write(self, text): 
 86          if text == None or text == "": 
 87              return 
 88          """ 
 89          Send the text to the stream after optionally writing it to the log file. 
 90          """ 
 91          if self.indent != None: 
 92              if self.newLine: 
 93                  text = self.indent + text  
 94              lastChar = text[-1] 
 95              text = text[:-1] 
 96              if lastChar == "\n": 
 97                  self.newLine = True 
 98              else: 
 99                  self.newLine = False 
100              text = text.replace("\n","\n"+self.indent) 
101              text += lastChar 
102          self.stream.write(text) 
103          self.stream.flush() 
104          if len(self.logfiles) > 0: 
105              for char in text: 
106                  if char == "\r": 
107                      self.buffer = "" 
108                  elif char == "\n": 
109                      timeString = time.strftime(self.timeStamp) 
110                      if timeString == self.prevTime and not self.timeStampDuplicates: 
111                          timeString = len(timeString) * " " 
112                      else: 
113                          self.prevTime = timeString 
114                      if self.timeStamp != None: 
115                          self.buffer = timeString + "\t" + self.buffer 
116                      for logfile in self.logfiles: 
117                          logfile.write(self.buffer + "\n") 
118                      self.buffer = "" 
119                  else: 
120                      self.buffer += char 
121              for logfile in self.logfiles: 
122                  logfile.flush() 
123       
124 -    def flush(self): 
125          self.stream.flush() 
126   
127 -def openLog(filename="log.txt", clear=False, logCmd=True): 
128      if os.path.dirname(filename) != "" and not os.path.exists(os.path.dirname(filename)): 
129          os.makedirs(os.path.dirname(filename)) 
130      setLog(filename, clear) 
131      setTimeStamp("[%H:%M:%S %d/%m]", True) 
132      logOpenTime = str(time.ctime(time.time())) 
133      print >> sys.stderr, "Opening log", filename, "at", logOpenTime 
134      logOpenMessage = "####### Log opened at " + str(time.ctime(time.time())) + " #######\n" 
135      writeToLog(logOpenMessage, filename) 
136      if logCmd: 
137          writeToLog("Command line: " + " ".join(sys.argv) + "\n", filename) 
138   
139 -def closeLog(filename): 
140      assert isinstance(sys.stdout, StreamModifier) 
141      removedStdout = sys.stdout.removeLog(filename, "stdout") 
142      assert isinstance(sys.stderr, StreamModifier) 
143      removedStderr = sys.stderr.removeLog(filename, "stderr") 
144      # These are most often the same file, so they (it) must be closed after removed from all streams 
145      removedStdout.close() 
146      removedStderr.close() 
147   
148 -def writeToScreen(text): 
149      assert isinstance(sys.stderr, StreamModifier) 
150      sys.stderr.writeToStream(text) 
151   
152 -def writeToLog(text, filename=None): 
153      assert isinstance(sys.stdout, StreamModifier) 
154      sys.stdout.writeToLog(text, filename) 
155   
156 -def setLog(filename=None, clear=False): 
157      """ 
158      Replace sys.stderr and sys.stdout with a StreamModifier, capturing 
159      all output for these streams to a log file while still passing it 
160      to the original stream. 
161      """ 
162      if not isinstance(sys.stdout, StreamModifier): 
163          sys.stdout = StreamModifier(sys.stdout) 
164      if not isinstance(sys.stderr, StreamModifier): 
165          sys.stderr = StreamModifier(sys.stderr) 
166      if filename != None: 
167          if clear: 
168              logfile = open(filename,"wt") 
169          else: 
170              logfile = open(filename,"at") 
171          sys.stdout.addLog(logfile) 
172          sys.stderr.addLog(logfile) 
173   
174 -def setIndent(string=None): 
175      if not isinstance(sys.stdout, StreamModifier): 
176          sys.stdout = StreamModifier(sys.stdout) 
177      if not isinstance(sys.stderr, StreamModifier): 
178          sys.stderr = StreamModifier(sys.stderr) 
179      sys.stdout.setIndent(string) 
180      sys.stderr.setIndent(string) 
181   
182 -def setTimeStamp(format=None, duplicates=False): 
183      if not isinstance(sys.stdout, StreamModifier): 
184          sys.stdout = StreamModifier(sys.stdout) 
185      if not isinstance(sys.stderr, StreamModifier): 
186          sys.stderr = StreamModifier(sys.stderr) 
187      sys.stdout.setTimeStamp(format, duplicates) 
188      sys.stderr.setTimeStamp(format, duplicates) 
189