data_logger.py

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