data_logger.py

   1 #!/user/bin/python
   2 #
   3 # Define some simple data logging classes for consistency
   4 # see the test functions for some usage examples
   5
   6 import os, os.path
   7 import stat
   8 import cPickle as pickle
   9 import time
  10 import string
  11 import numpy
  12
  13 class error (Exception) :
  14     "Basic module error class"
  15     pass
  16
  17 class errorDirExists (error) :
  18     "The specified directory already exists"
  19
  20 class data_log :
  21     """
  22     Data logging class.
  23     Creates consistent, timestamped log files.
  24
  25     Initialized with log_dir and log_name.
  26     log_dir specifies the base data directory.
  27     If it doesn't exist, log_dir is created.
  28
  29     A subdir of log_dir is created (if necessary) named YYYYMMDD,
  30     where YYYYMMDD is the current day in localtime.
  31     If noclobber_logsubdir == True, this dir must not exist yet.
  32
  33     log_name specifies the base name for the created log files (in the log subdir).
  34     The created log filenames are prefixed with a YYYYMMDDHHMMSS timestamp.
  35     If the target filename already exists, the filename is postfixed with
  36     '_N', where N is the lowest integer that doesn't clobber an existing file.
  37
  38     General data is saved to the log files with the write(obj) method.
  39     By default, write() cPickles the object passed.
  40     You can save in other formats by overriding write()
  41
  42     Binary data is can be saved directly to the log files with the
  43     write_binary(binary_string) method.
  44
  45     All file names are stripped of possibly troublesome characters.
  46     """
  47     def __init__(self, log_dir=".", noclobber_logsubdir=False,
  48                  log_name="log",
  49                  timestamp=None) :
  50         # generate lists of not-allowed characters
  51         unaltered_chars = "-._" + string.digits + string.letters
  52         mapped_pairs = {' ':'_'}
  53         allowed_chars = unaltered_chars + "".join(mapped_pairs.keys())
  54         all_chars = string.maketrans('','')
  55         self.delete_chars = all_chars.translate(all_chars, allowed_chars)
  56         trans_from = "".join(mapped_pairs.keys())
  57         trans_to = "".join(mapped_pairs.values()) # same order as keys, since no modifications to mapped_pairs were made in between the two calls
  58         self.transtable = string.maketrans(trans_from, trans_to)
  59
  60         self._log_name = self._clean_filename(log_name) # never checked after this...
  61         self._log_dir = self._create_logdir(log_dir) # will not clobber.
  62         subdir, timestamp = self._create_logsubdir(self._log_dir,
  63                                                    noclobber_logsubdir,
  64                                                    timestamp)
  65         self.subdir = subdir
  66         self.timestamp = timestamp
  67     def _clean_filename(self, filename) :
  68         """
  69         Currently only works on filenames, since it deletes '/'.
  70         If you need it to work on full paths, use os.path.split(your_path)[1]
  71         to strip of the filename portion...
  72         """
  73         cleanname = filename.translate(self.transtable, self.delete_chars)
  74         return cleanname
  75     def _create_logdir(self, log_dir) :
  76         log_dir = os.path.expanduser(log_dir)
  77         if not os.path.exists(log_dir) :
  78             os.mkdir(log_dir, 0755)
  79         return log_dir
  80     def _create_logsubdir(self, log_dir, noclobber_logsubdir,
  81                           timestamp=None) :
  82         if timestamp == None :
  83             timestamp = time.strftime("%Y%m%d") # %H%M%S
  84         subdir = os.path.join(log_dir, timestamp)
  85         if os.path.exists(subdir) :
  86             if noclobber_logsubdir:
  87                 raise errorDirExists, "%s exists" % subdir
  88         else :
  89             os.mkdir(subdir, 0755)
  90         return (subdir, timestamp)
  91     def get_filename(self, timestamp=None) :
  92         """
  93         Get a filename (using localtime if timestamp==None),
  94         appending integers as necessary to avoid clobbering.
  95         For use in write() routines.
  96         Returns (filepath, timestamp)
  97         """
  98         if timestamp == None :
  99             timestamp = time.strftime("%Y%m%d%H%M%S")
 100         filename = "%s_%s" % (timestamp, self._log_name)
 101         fullname = os.path.join(self.subdir, filename)
 102         filepath = fullname
 103         i = 1
 104         while os.path.exists(filepath) :
 105             filepath = "%s_%d" % (fullname, i)
 106             i+=1
 107         return (filepath, timestamp)
 108     def write(self, obj, timestamp=None) :
 109         """
 110         Save object to a timestamped file with pickle.
 111         If timestamp == None, use the current localtime.
 112         Returns (filepath, timestamp)
 113         """
 114         filepath, timestamp = self.get_filename(timestamp)
 115         fd = open(filepath, 'wb')
 116         os.chmod(filepath, 0644)
 117         pickle.dump(obj, fd)
 118         fd.close()
 119         return (filepath, timestamp)
 120     def write_binary(self, binary_string, timestamp=None) :
 121         """
 122         Save binary_string to a timestamped file.
 123         If timestamp == None, use the current localtime.
 124         Returns (filepath, timestamp)
 125         """
 126         filepath, timestamp = self.get_filename(timestamp)
 127         # open a new file in readonly mode, don't clobber.
 128         fd = os.open(filepath, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0644)
 129         bytes_written = 0
 130         bytes_remaining = len(binary_string)
 131         while bytes_remaining > 0 :
 132             bw = os.write(fd, binary_string[bytes_written:])
 133             bytes_written += bw
 134             bytes_remaining -= bw
 135         os.close(fd)
 136         return (filepath, timestamp)
 137     def _write_dict_of_arrays(self, d, base_filepath) :
 138         # open a new file in readonly mode, don't clobber.
 139         bfd = open(base_filepath, 'w', 0644)
 140         bfd.write("Contents (key : file-extension : format):\n")
 141         for key in d.keys() :
 142             clean_key = self._clean_filename(key)
 143             bfd.write("%s : %s : %s\n" % (key, clean_key, str(d[key].dtype)))
 144             # write the keyed array to it's own file
 145             filepath = "%s_%s" % (base_filepath, clean_key)
 146             d[key].tofile(filepath)
 147         bfd.close()
 148     def write_dict_of_arrays(self, d, timestamp=None) :
 149         """
 150         Save dict of (string, numpy_array) pairs to timestamped files.
 151         If timestamp == None, use the current localtime.
 152         Returns (base_filepath, timestamp)
 153         """
 154         base_filepath, timestamp = self.get_filename(timestamp)
 155         self._write_dict_of_arrays(d, base_filepath)
 156         return (base_filepath, timestamp)
 157
 158 class data_load :
 159     """
 160     Loads data logged by data_log.
 161     """
 162     def read(self, file) :
 163         """
 164         Load an object saved with data_log.write()
 165         """
 166         return pickle.load(open(file, 'rb'))
 167     def read_binary(self, file) :
 168         """
 169         Load an object saved with data_log.write_binary()
 170         The file-name must not have been altered.
 171         """
 172         raise Exception, "not implemented"
 173     def read_dict_of_arrays(self, basefile) :
 174         """
 175         Load an object saved with data_log.write_binary()
 176         The file-names must not have been altered.
 177         """
 178         obj = {}
 179         i=0
 180         realbasefile = os.path.realpath(basefile)
 181         for line in file(realbasefile) :
 182             if i > 0 : # ignore first line
 183                 ldata = line.split(' : ')
 184                 name = ldata[0]
 185                 fpath = "%s_%s" % (realbasefile, ldata[1])
 186                 exec 'typ = numpy.%s' % ldata[2]
 187                 obj[name] = numpy.fromfile(fpath, dtype=typ)
 188             i += 1
 189         return obj
 190
 191 _test_dir = "."
 192
 193 def _check_data_logsubdir_clobber() :
 194     log1 = data_log(_test_dir, noclobber_logsubdir=True)
 195     try :
 196         log2 = data_log(_test_dir, noclobber_logsubdir=True)
 197         raise error, "Didn't detect old log"
 198     except errorDirExists :
 199         pass # everything as it should be
 200     os.rmdir(log1.subdir)
 201
 202 def _check_data_log_filenames() :
 203     data = {"Test":True, "Data":[1,2,3,4]}
 204     log = data_log(_test_dir, noclobber_logsubdir=True)
 205     files = [None]*10
 206     for i in range(10):
 207         files[i], ts = log.write(data)
 208     print "Contents of log directory (should be 10 identical logs)"
 209     os.system('ls -l %s' % log.subdir)
 210     for file in files :
 211         os.remove( file )
 212     os.rmdir(log.subdir)
 213
 214 def _check_data_log_pickle_integrity() :
 215     data = {"Test":True, "Data":[1,2,3,4]}
 216     # save the data
 217     log = data_log(_test_dir, noclobber_logsubdir=True)
 218     filepath, ts = log.write(data)
 219     # read it back in
 220     fd = open(filepath, 'rb')
 221     data_in = pickle.load(fd)
 222     fd.close()
 223     # compare
 224     if data != data_in :
 225         print "Saved    : ", data
 226         print "Read back: ", data_in
 227         raise error, "Poorly pickled"
 228     os.remove(filepath)
 229     os.rmdir(log.subdir)
 230
 231 def _check_data_log_binary_integrity() :
 232     from numpy import zeros, uint16, fromfile
 233     npts = 100
 234     data = zeros((npts,), dtype=uint16)
 235     for i in range(npts) :
 236         data[i] = i
 237     # save the data
 238     log = data_log(_test_dir, noclobber_logsubdir=True)
 239     filepath, ts = log.write_binary(data.tostring())
 240     # read it back in
 241     data_in = fromfile(filepath, dtype=uint16, count=-1)
 242     # compare
 243     if npts != len(data_in) :
 244         raise error, "Saved %d uint16s, read %d" % (npts, len(data_in))
 245     for i in range(npts) :
 246         if data_in[i] != data[i] :
 247             print "Disagreement in element %d" % i
 248             print "Saved %d, read back %d" % (data[i], data_in[i])
 249             raise error, "Poorly saved"
 250     os.remove(filepath)
 251     os.rmdir(log.subdir)
 252
 253 def _check_data_loc_dict_of_arrays() :
 254     from numpy import zeros, uint16, fromfile
 255     npts = 100
 256     data1 = zeros((npts,), dtype=uint16)
 257     for i in range(npts) :
 258         data1[i] = i
 259     data2 = zeros((npts,), dtype=uint16)
 260     for i in range(npts) :
 261         data2[i] = npts-i
 262     data={"data1":data1, "d\/at:$a 2":data2}
 263     # save the data
 264     log = data_log(_test_dir, noclobber_logsubdir=True)
 265     filepath, ts = log.write_dict_of_arrays(data)
 266     # checking
 267     print "Contents of log directory (should be 3 logs)"
 268     os.system('ls -l %s' % log.subdir)
 269     print "The table of contents file:"
 270     os.system('cat %s' % (filepath))
 271     data1_in = fromfile(filepath+"_data1", dtype=uint16)
 272     data2_in = fromfile(filepath+"_data_2", dtype=uint16)
 273     for i in range(npts) :
 274         if data1_in[i] != data1[i] :
 275             print "Disagreement in element %d of data1" % i
 276             print "Saved %d, read back %d" % (data1[i], data1_in[i])
 277             raise error, "Poorly saved"
 278         if data2_in[i] != data2[i] :
 279             print "Disagreement in element %d of data2" % i
 280             print "Saved %d, read back %d" % (data2[i], data2_in[i])
 281             raise error, "Poorly saved"
 282     os.remove(filepath)
 283     os.remove(filepath+"_data1")
 284     os.remove(filepath+"_data_2")
 285     os.rmdir(log.subdir)
 286
 287 def test() :
 288     _check_data_logsubdir_clobber()
 289     _check_data_log_filenames()
 290     _check_data_log_pickle_integrity()
 291     _check_data_log_binary_integrity()
 292     _check_data_loc_dict_of_arrays()
 293
 294 if __name__ == "__main__" :
 295     test()