Source Code for Module io

  1  ############################################################################### 
  2  #                                                                             # 
  3  # Copyright (C) 2003-2005 Edward d'Auvergne                                   # 
  4  #                                                                             # 
  5  # This file is part of the program relax.                                     # 
  6  #                                                                             # 
  7  # relax is free software; you can redistribute it and/or modify               # 
  8  # it under the terms of the GNU General Public License as published by        # 
  9  # the Free Software Foundation; either version 2 of the License, or           # 
 10  # (at your option) any later version.                                         # 
 11  #                                                                             # 
 12  # relax is distributed in the hope that it will be useful,                    # 
 13  # but WITHOUT ANY WARRANTY; without even the implied warranty of              # 
 14  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               # 
 15  # GNU General Public License for more details.                                # 
 16  #                                                                             # 
 17  # You should have received a copy of the GNU General Public License           # 
 18  # along with relax; if not, write to the Free Software                        # 
 19  # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA   # 
 20  #                                                                             # 
 21  ############################################################################### 
 22   
 23  import __builtin__ 
 24   
 25  # BZ2 compression module. 
 26  try: 
 27      from bz2 import BZ2File 
 28      __builtin__.bz2_module = 1 
 29  except ImportError, message: 
 30      __builtin__.bz2_module = 0 
 31      __builtin__.bz2_module_message = message.args[0] 
 32   
 33  # Gzip compression module. 
 34  from gzip import GzipFile 
 35   
 36  from os import F_OK, access, makedirs, remove, stat 
 37  from os.path import expanduser 
 38  from re import match, search 
 39  from string import split 
 40  import sys 
 41  from sys import stdin, stdout, stderr 
 42   
 43   
 44 -class IO: 
 45 -    def __init__(self, relax): 
 46          """Class containing the file operations. 
 47   
 48          IO streams 
 49          ~~~~~~~~~~ 
 50   
 51          Standard python IO streams: 
 52   
 53          sys.stdin  = self.python_stdin 
 54          sys.stdout = self.python_stdout 
 55          sys.stderr = self.python_stderr 
 56   
 57          Logging IO streams: 
 58   
 59          sys.stdin  = self.log_stdin  = self.python_stdin 
 60          sys.stdout = self.log_stdout = self.log_file 
 61          sys.stderr = self.log_stdout = (self.python_stderr, self.log_file) 
 62   
 63          Tee IO streams: 
 64   
 65          sys.stdin  = self.tee_stdin  = self.python_stdin 
 66          sys.stdout = self.tee_stdout = (self.python_stdout, self.tee_file) 
 67          sys.stderr = self.tee_stdout = (self.python_stderr, self.tee_file) 
 68          """ 
 69   
 70          self.relax = relax 
 71   
 72          # Standard python IO streams. 
 73          self.python_stdin  = stdin 
 74          self.python_stdout = stdout 
 75          self.python_stderr = stderr 
 76   
 77          # Logging IO streams. 
 78          self.log_stdin  = stdin 
 79          self.log_stdout = None 
 80          self.log_stderr = SplitIO() 
 81   
 82          # Tee IO streams. 
 83          self.tee_stdin  = stdin 
 84          self.tee_stdout = SplitIO() 
 85          self.tee_stderr = SplitIO() 
 86   
 87   
 88 -    def delete(self, file_name=None, dir=None): 
 89          """Function for deleting the given file.""" 
 90   
 91          # File path. 
 92          file_path = self.file_path(file_name, dir) 
 93   
 94          # Test if the file exists and determine the compression type. 
 95          if access(file_path, F_OK): 
 96              pass 
 97          elif access(file_path + '.bz2', F_OK): 
 98              file_path = file_path + '.bz2' 
 99          elif access(file_path + '.gz', F_OK): 
100              file_path = file_path + '.gz' 
101          else: 
102              raise RelaxFileError, file_path 
103   
104          # Remove the file. 
105          remove(file_path) 
106   
107   
108 -    def extract_data(self, file_name=None, dir=None, file_data=None, sep=None, compress_type=0): 
109          """Open the file 'file' and return all the data.""" 
110   
111          # Data not already extracted from the file. 
112          if not file_data: 
113              # Open the file. 
114              file = self.open_read_file(file_name=file_name, dir=dir, compress_type=compress_type) 
115   
116              # Read lines. 
117              file_data = file.readlines() 
118   
119          # Create a data structure from the contents of the file split by either whitespace or the separator, sep. 
120          data = [] 
121          for i in xrange(len(file_data)): 
122              if sep: 
123                  row = split(file_data[i], sep) 
124              else: 
125                  row = split(file_data[i]) 
126              data.append(row) 
127          return data 
128   
129          # Close the file. 
130          if not file_data: 
131              file.close() 
132   
133   
134 -    def file_path(self, file_name=None, dir=None): 
135          """Generate and expand the full file path.""" 
136   
137          # File name. 
138          file_path = file_name 
139   
140          # Add the directory. 
141          if dir: 
142              file_path = dir + '/' + file_path 
143   
144          # Expand any ~ characters. 
145          file_path = expanduser(file_path) 
146   
147          # Return the file path. 
148          return file_path 
149   
150   
151 -    def log(self, file_name=None, dir=None, compress_type=0, print_flag=1): 
152          """Function for turning logging on.""" 
153   
154          # Log file. 
155          self.log_file, file_path = self.open_write_file(file_name=file_name, dir=dir, force=1, compress_type=compress_type, print_flag=print_flag, return_path=1) 
156   
157          # Print out. 
158          if print_flag: 
159              print "Redirecting the sys.stdin IO stream to the python stdin IO stream." 
160              print "Redirecting the sys.stdout IO stream to the log file '%s'." % file_path 
161              print "Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path 
162   
163          # Set the logging IO streams. 
164          self.log_stdout = self.log_file 
165          self.log_stderr.split(self.python_stderr, self.log_file) 
166   
167          # IO stream redirection. 
168          sys.stdin  = self.log_stdin 
169          sys.stdout = self.log_stdout 
170          sys.stderr = self.log_stderr 
171   
172   
173 -    def logging_off(self, file_name=None, dir=None, print_flag=1): 
174          """Function for turning logging and teeing off.""" 
175   
176          # Print out. 
177          if print_flag: 
178              print "Redirecting the sys.stdin IO stream to the python stdin IO stream." 
179              print "Redirecting the sys.stdout IO stream to the python stdout IO stream." 
180              print "Redirecting the sys.stderr IO stream to the python stderr IO stream." 
181   
182          # IO stream redirection. 
183          sys.stdin  = self.python_stdin 
184          sys.stdout = self.python_stdout 
185          sys.stderr = self.python_stderr 
186   
187   
188 -    def mkdir(self, dir=None, print_flag=1): 
189          """Create the given directory, or exit if the directory exists.""" 
190   
191          # No directory given. 
192          if dir == None: 
193              return 
194   
195          # Make the directory. 
196          try: 
197              makedirs(dir) 
198          except OSError: 
199              if print_flag: 
200                  print "Directory ./" + dir + " already exists.\n" 
201   
202   
203 -    def open_read_file(self, file_name=None, dir=None, compress_type=0, print_flag=1): 
204          """Open the file 'file' and return all the data.""" 
205   
206          # File path. 
207          file_path = self.file_path(file_name, dir) 
208   
209          # Test if the file exists and determine the compression type. 
210          if access(file_path, F_OK): 
211              compress_type = 0 
212              if search('.bz2$', file_path): 
213                  compress_type = 1 
214              elif search('.gz$', file_path): 
215                  compress_type = 2 
216          elif access(file_path + '.bz2', F_OK): 
217              file_path = file_path + '.bz2' 
218              compress_type = 1 
219          elif access(file_path + '.gz', F_OK): 
220              file_path = file_path + '.gz' 
221              compress_type = 2 
222          else: 
223              raise RelaxFileError, file_path 
224   
225          # Open the file for reading. 
226          try: 
227              if print_flag: 
228                  print "Opening the file " + `file_path` + " for reading." 
229              if compress_type == 0: 
230                  file = open(file_path, 'r') 
231              elif compress_type == 1: 
232                  if bz2_module: 
233                      file = BZ2File(file_path, 'r') 
234                  else: 
235                      raise RelaxError, "Cannot open the file " + `file_path` + ", try uncompressing first.  " + bz2_module_message + "." 
236              elif compress_type == 2: 
237                  file = GzipFile(file_path, 'r') 
238          except IOError, message: 
239              raise RelaxError, "Cannot open the file " + `file_path` + ".  " + message.args[1] + "." 
240   
241          # Return the opened file. 
242          return file 
243   
244   
245 -    def open_write_file(self, file_name=None, dir=None, force=0, compress_type=0, print_flag=1, return_path=0): 
246          """Function for opening a file for writing and creating directories if necessary.""" 
247   
248          # Create the directories. 
249          self.mkdir(dir, print_flag=0) 
250   
251          # File path. 
252          file_path = self.file_path(file_name, dir) 
253   
254          # Bzip2 compression. 
255          if compress_type == 1 and not search('.bz2$', file_path): 
256              # Bz2 module exists. 
257              if bz2_module: 
258                  file_path = file_path + '.bz2' 
259   
260              # Switch to gzip compression. 
261              else: 
262                  print "Cannot use bz2 compression, using gzip compression instead.  " + bz2_module_message + "." 
263                  compress_type = 2 
264   
265          # Gzip compression. 
266          if compress_type == 2 and not search('.gz$', file_path): 
267              file_path = file_path + '.gz' 
268   
269          # Fail if the file already exists and the force flag is set to 0. 
270          if access(file_path, F_OK) and not force: 
271              raise RelaxFileOverwriteError, (file_path, 'force flag') 
272   
273          # Open the file for writing. 
274          try: 
275              if print_flag: 
276                  print "Opening the file " + `file_path` + " for writing." 
277              if compress_type == 0: 
278                  file = open(file_path, 'w') 
279              elif compress_type == 1: 
280                  file = BZ2File(file_path, 'w') 
281              elif compress_type == 2: 
282                  file = GzipFile(file_path, 'w') 
283          except IOError, message: 
284              raise RelaxError, "Cannot open the file " + `file_path` + ".  " + message.args[1] + "." 
285   
286          # Return the opened file. 
287          if return_path: 
288              return file, file_path 
289          else: 
290              return file 
291   
292   
293 -    def strip(self, data): 
294          """Function to remove all comment and empty lines from the file data structure.""" 
295   
296          # Initialise the new data array. 
297          new = [] 
298   
299          # Loop over the data. 
300          for i in xrange(len(data)): 
301              # Empty lines. 
302              if len(data[i]) == 0: 
303                  continue 
304   
305              # Comment lines. 
306              elif match("#", data[i][0]): 
307                  continue 
308   
309              # Data lines. 
310              else: 
311                  new.append(data[i]) 
312   
313          # Return the new data structure. 
314          return new 
315   
316   
317 -    def tee(self, file_name=None, dir=None, compress_type=0, print_flag=1): 
318          """Function for turning logging on.""" 
319   
320          # Tee file. 
321          self.tee_file, file_path = self.open_write_file(file_name=file_name, dir=dir, force=1, compress_type=compress_type, print_flag=print_flag, return_path=1) 
322   
323          # Print out. 
324          if print_flag: 
325              print "Redirecting the sys.stdin IO stream to the python stdin IO stream." 
326              print "Redirecting the sys.stdout IO stream to both the python stdout IO stream and the log file '%s'." % file_path 
327              print "Redirecting the sys.stderr IO stream to both the python stderr IO stream and the log file '%s'." % file_path 
328   
329          # Set the tee IO streams. 
330          self.tee_stdout.split(self.python_stdout, self.tee_file) 
331          self.tee_stderr.split(self.python_stderr, self.tee_file) 
332   
333          # IO stream redirection. 
334          sys.stdin  = self.tee_stdin 
335          sys.stdout = self.tee_stdout 
336          sys.stderr = self.tee_stderr 
337   
338   
339   
340 -class SplitIO: 
341 -    def __init__(self): 
342          """Class for splitting an IO stream to two outputs.""" 
343   
344   
345 -    def split(self, stream1, stream2): 
346          """Function for setting the streams.""" 
347   
348          # Arguments. 
349          self.stream1 = stream1 
350          self.stream2 = stream2 
351   
352   
353 -    def write(self, text): 
354          """Replacement write function.""" 
355   
356          # Write to stream1. 
357          self.stream1.write(text) 
358   
359          # Write to stream2. 
360          self.stream2.write(text) 
361