data_logger.py

   1 #!/user/bin/python
   2 #
   3 # data_logger - classes for consistently logging data in an organized
   4 # fasion.  See the doctests for some usage examples.
   5 #
   6 # Copyright (C) 2008-2010 William Trevor King
   7 #
   8 # This program is free software; you can redistribute it and/or
   9 # modify it under the terms of the GNU General Public License as
  10 # published by the Free Software Foundation; either version 3 of the
  11 # License, or (at your option) any later version.
  12 #
  13 # This program is distributed in the hope that it will be useful, but
  14 # WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
  16 # See the GNU General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with this program; if not, write to the Free Software
  20 # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA
  21 # 02111-1307, USA.
  22 #
  23 # The author may be contacted at <wking@drexel.edu> on the Internet, or
  24 # write to Trevor King, Drexel University, Physics Dept., 3141 Chestnut St.,
  25 # Philadelphia PA 19104, USA.
  26
  27 from __future__ import with_statement
  28
  29 import os, os.path
  30 import stat
  31 import cPickle as pickle
  32 import time
  33 import string
  34
  35 import numpy
  36
  37
  38 __version__ = "0.4"
  39 DEFAULT_PATH = "~/rsrch/data"
  40 DEFAULT_PATH_REPLACE_STRING = "${DEFAULT}/"
  41
  42
  43 class Error (Exception):
  44     "Basic module error class"
  45     pass
  46
  47
  48 class ErrorDirExists (Error):
  49     "The specified directory already exists"
  50     pass
  51
  52
  53 def normalize_log_dir(log_dir):
  54     """Normalize a log directory.
  55
  56     Expands the user symbol `~`, as well as
  57     `DEFAULT_PATH_REPLACE_STRING`.
  58
  59     Parameters
  60     ----------
  61     log_dir : path
  62         Raw `log_dir` passed into `.__init__()`.
  63
  64     Returns
  65     -------
  66     log_dir : path
  67         Normalized version of the input `log_dir`.
  68
  69     Examples
  70     --------
  71     >>> normalize_log_dir('~/.log')  # doctest: +ELLIPSIS
  72     '/.../.log'
  73     >>> normalize_log_dir('${DEFAULT}/hi/there')  # doctest: +ELLIPSIS
  74     '/.../rsrch/data/hi/there'
  75     """
  76     if log_dir.startswith(DEFAULT_PATH_REPLACE_STRING):
  77         length = len(DEFAULT_PATH_REPLACE_STRING)
  78         log_dir = os.path.join(DEFAULT_PATH, log_dir[length:])
  79     log_dir = os.path.expanduser(log_dir)
  80     return log_dir
  81
  82
  83 class DataLog (object):
  84     """Create consistent, timestamped log files.
  85
  86     General data is saved to the log files with the `write(obj)`
  87     method.  By default, `write()` `cPickles` the object passed.  You
  88     can save in other formats by overriding `write()`.
  89
  90     Binary data is can be saved directly to the log files with the
  91     `write_binary(binary_string)` method.
  92
  93     All file names are stripped of possibly troublesome characters.
  94
  95     Parameters
  96     ----------
  97     log_dir : path
  98         `log_dir` sets the base data directory.  If it doesn't exist,
  99         `log_dir` is created.
 100
 101         If log_dir begins with '${DEFAULT}/', that portion of the path
 102         is replaced with the then-current contents of the
 103         `DEFAULT_PATH` module global.
 104
 105         A subdir of log_dir is created (if necessary) named
 106         `YYYYMMDD`, where `YYYYMMDD` is the current day in local time.
 107     log_name : string
 108          `log_name` specifies the base name for the created log files
 109          (in the log subdir).  The created log filenames are prefixed
 110          with a `YYYYMMDDHHMMSS` timestamp.  If the target filename
 111          already exists, the filename is postfixed with `_N`, where
 112          `N` is the lowest integer that doesn't clobber an existing
 113          file.
 114     noclobber_log_subdir : bool
 115         `noclobber_log_subdir == True`, the `YYYMMDD` subdir of
 116         `log_dir` must not exist yet.
 117     timestamp : string
 118          Overide default subdir `timestamp` (%Y%m%d).
 119
 120     Examples
 121     --------
 122
 123     >>> import shutil
 124     >>> dl = DataLog('test_data_log', 'temperature', timestamp='20101103',
 125     ...              )
 126     >>> data = {'test':True, 'data':[1, 2, 3, 4]}
 127     >>> files = [None]*10
 128     >>> for i in range(10):
 129     ...     files[i],ts = dl.write(data, timestamp='20101103235959')
 130     >>> print '\\n'.join(files)
 131     test_data_log/20101103/20101103235959_log
 132     test_data_log/20101103/20101103235959_log_1
 133     test_data_log/20101103/20101103235959_log_2
 134     test_data_log/20101103/20101103235959_log_3
 135     test_data_log/20101103/20101103235959_log_4
 136     test_data_log/20101103/20101103235959_log_5
 137     test_data_log/20101103/20101103235959_log_6
 138     test_data_log/20101103/20101103235959_log_7
 139     test_data_log/20101103/20101103235959_log_8
 140     test_data_log/20101103/20101103235959_log_9
 141     >>> shutil.rmtree(dl._log_dir)
 142     """
 143     def __init__(self, log_dir=".", noclobber_log_subdir=False,
 144                  log_name="log", timestamp=None):
 145         self._setup_character_translation()
 146         self._log_name = self._clean_filename(log_name)  # last check.
 147         self._log_dir = self._create_log_dir(log_dir)  # will not clobber.
 148         self._subdir,self._timestamp = self._create_log_subdir(
 149             self._log_dir, noclobber_log_subdir, timestamp)
 150
 151     def _setup_character_translation(self):
 152         """Setup `._delete_chars` and `._trans_table` for `._clean_filename()`.
 153         """
 154         # generate lists of not-allowed characters
 155         unaltered_chars = '-._' + string.digits + string.letters
 156         mapped_pairs = {' ':'_'}
 157         allowed_chars = unaltered_chars + ''.join(mapped_pairs.keys())
 158         all_chars = string.maketrans('','')
 159         self._delete_chars = all_chars.translate(all_chars, allowed_chars)
 160         trans_from = ''.join(mapped_pairs.keys())
 161         trans_to = ''.join(mapped_pairs.values())
 162         # values in trans_to are in the same order as the keys in
 163         # trans_from, since no modifications to mapped_pairs were made
 164         # in between the two calls.
 165         self._trans_table = string.maketrans(trans_from, trans_to)
 166
 167     def _clean_filename(self, filename):
 168         """Remove troublesome characters from filenames.
 169
 170         This method only works on filenames, since it deletes '/'.  If
 171         you need it to work on full paths, use
 172         `os.path.split(your_path)` and clean the portions separately.
 173
 174         Parameters
 175         ----------
 176         filename : string
 177
 178         Examples
 179         --------
 180         >>> import shutil
 181         >>> dl = DataLog(log_dir="test_clean_filename")
 182         >>> dl._clean_filename('hi there')
 183         'hi_there'
 184         >>> dl._clean_filename('hello\\tthe/castle')
 185         'hellothecastle'
 186         >>> shutil.rmtree(dl._log_dir)
 187         """
 188         cleanname = filename.translate(self._trans_table, self._delete_chars)
 189         return cleanname
 190
 191     def _create_log_dir(self, log_dir):
 192         """Create a clean base log dir (if necessary).
 193
 194         Parameters
 195         ----------
 196         log_dir : path
 197             Raw `log_dir` passed into `.__init__()`.
 198
 199         Returns
 200         -------
 201         log_dir : path
 202             Normalized version of the input `log_dir`.
 203
 204         Examples
 205         --------
 206         >>> import shutil
 207         >>> dl = DataLog(log_dir='test_create_log_dir')
 208         >>> shutil.rmtree(dl._log_dir)
 209         """
 210         log_dir = normalize_log_dir(log_dir)
 211         if not os.path.exists(log_dir):
 212             os.mkdir(log_dir, 0755)
 213         return log_dir
 214
 215     def _create_log_subdir(self, log_dir, noclobber_log_subdir=False,
 216                           timestamp=None):
 217         """Create a clean log dir for logging.
 218
 219         Parameters
 220         ----------
 221         log_dir : path
 222             Normalized version of the input `log_dir`.
 223         noclobber_log_subdir : bool
 224             `noclobber_log_subdir` passed into `.__init__()`.
 225         timestamp : string
 226             Overide default `timestamp` (%Y%m%d).
 227
 228         Returns
 229         -------
 230         subdir : path
 231             Path to the timestamped subdir of `log_dir`.
 232         timestamp : string
 233             The timestamp used to generate `subdir`.
 234
 235         Examples
 236         --------
 237         >>> import os
 238         >>> import shutil
 239         >>> dl = DataLog(log_dir='test_create_log_subdir',
 240         ...              timestamp='20101103')
 241         >>> os.listdir(dl._log_dir)
 242         ['20101103']
 243         >>> dl._create_log_subdir(dl._log_dir, noclobber_log_subdir=True,
 244         ...                       timestamp=dl._timestamp)
 245         Traceback (most recent call last):
 246           ...
 247         ErrorDirExists: test_create_log_subdir/20101103 exists
 248         >>> dl._create_log_subdir(dl._log_dir, noclobber_log_subdir=False,
 249         ...                       timestamp=dl._timestamp)
 250         ('test_create_log_subdir/20101103', '20101103')
 251         >>> dl._create_log_subdir(dl._log_dir)  # doctest: +ELLIPSIS
 252         ('test_create_log_subdir/...', '...')
 253         >>> shutil.rmtree(dl._log_dir)
 254         """
 255         if timestamp == None:
 256             timestamp = time.strftime("%Y%m%d") # %H%M%S
 257         subdir = os.path.join(log_dir, timestamp)
 258         if os.path.exists(subdir):
 259             if noclobber_log_subdir:
 260                 raise ErrorDirExists, "%s exists" % subdir
 261         else:
 262             os.mkdir(subdir, 0755)
 263         return (subdir, timestamp)
 264
 265     def _get_filename(self, timestamp=None):
 266         """Get a filename for a new data log for `.write()`.
 267
 268         Append integers as necessary to avoid clobbering.  Note that
 269         the appended integers are *not* thread-safe.  You need to
 270         actually create the file to reserve the name.
 271
 272         Parameters
 273         ----------
 274         log_dir : path
 275             Normalized version of the input `log_dir`.
 276         noclobber_log_subdir : bool
 277             `noclobber_log_subdir` passed into `.__init__()`.
 278         timestamp : string
 279             Overide default `timestamp` (%Y%m%d%H%M%S).
 280
 281         Returns
 282         -------
 283         filepath : path
 284             Path to the timestamped log file.
 285         timestamp : string
 286             The timestamp used to generate `subdir`.
 287
 288         Examples
 289         --------
 290         >>> import shutil
 291         >>> dl = DataLog(log_dir='test_get_filename',
 292         ...              log_name='my-log', timestamp='20101103')
 293         >>> f,t = dl._get_filename('20100103235959')
 294         >>> f
 295         'test_get_filename/20101103/20100103235959_my-log'
 296         >>> t
 297         '20100103235959'
 298         >>> open(f, 'w').write('dummy content')
 299         >>> f,t = dl._get_filename('20100103235959')
 300         >>> f
 301         'test_get_filename/20101103/20100103235959_my-log_1'
 302         >>> t
 303         '20100103235959'
 304         >>> open(f, 'w').write('dummy content')
 305         >>> f,t = dl._get_filename('20100103235959')
 306         >>> f
 307         'test_get_filename/20101103/20100103235959_my-log_2'
 308         >>> t
 309         '20100103235959'
 310         >>> dl._get_filename()  # doctest: +ELLIPSIS
 311         ('test_get_filename/20101103/..._my-log', '...')
 312         >>> shutil.rmtree(dl._log_dir)
 313         """
 314         if timestamp == None:
 315             timestamp = time.strftime("%Y%m%d%H%M%S")
 316         filename = "%s_%s" % (timestamp, self._log_name)
 317         fullname = os.path.join(self._subdir, filename)
 318         filepath = fullname
 319         i = 1
 320         while os.path.exists(filepath):
 321             filepath = "%s_%d" % (fullname, i)
 322             i+=1
 323         return (filepath, timestamp)
 324
 325     def write(self, obj, timestamp=None):
 326         """Save object to a timestamped file with `cPickle`.
 327
 328         Parameters
 329         ----------
 330         obj : object
 331             Object to save.
 332         timestamp : string
 333             Passed on to `._get_filename()`.
 334
 335         Returns
 336         -------
 337         filepath : path
 338             Path to the timestamped log file.
 339         timestamp : string
 340             The timestamp used to generate the log file.
 341
 342         Examples
 343         --------
 344         >>> import shutil
 345         >>> dl = DataLog(log_dir='test_write',
 346         ...              log_name='my-log', timestamp='20101103')
 347         >>> f,t = dl.write([1, 2, 3])
 348         >>> a = pickle.load(open(f, 'rb'))
 349         >>> a
 350         [1, 2, 3]
 351         >>> shutil.rmtree(dl._log_dir)
 352         """
 353         filepath, timestamp = self._get_filename(timestamp)
 354         with open(filepath, 'wb') as fd:
 355             os.chmod(filepath, 0644)
 356             pickle.dump(obj, fd)
 357         return (filepath, timestamp)
 358
 359     def write_binary(self, binary_string, timestamp=None):
 360         """Save a binary string to a timestamped file.
 361
 362         Parameters
 363         ----------
 364         binary_string : buffer
 365             Binary string to save.
 366         timestamp : string
 367             Passed on to `._get_filename()`.
 368
 369         Returns
 370         -------
 371         filepath : path
 372             Path to the timestamped log file.
 373         timestamp : string
 374             The timestamp used to generate the log file(s).
 375
 376         Examples
 377         --------
 378         >>> import shutil
 379         >>> import numpy
 380         >>> dl = DataLog(log_dir='test_write_binary',
 381         ...              log_name='my-log', timestamp='20101103')
 382         >>> data = numpy.arange(5, dtype=numpy.uint16)
 383         >>> filepath,ts = dl.write_binary(data.tostring())
 384         >>> data_in = numpy.fromfile(filepath, dtype=numpy.uint16, count=-1)
 385         >>> data_in
 386         array([0, 1, 2, 3, 4], dtype=uint16)
 387         >>> (data == data_in).all()
 388         True
 389         >>> shutil.rmtree(dl._log_dir)
 390         """
 391         filepath, timestamp = self._get_filename(timestamp)
 392         # open a new file in readonly mode, don't clobber.
 393         fd = os.open(filepath, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0644)
 394         bytes_written = 0
 395         bytes_remaining = len(binary_string)
 396         while bytes_remaining > 0:
 397             bw = os.write(fd, binary_string[bytes_written:])
 398             bytes_written += bw
 399             bytes_remaining -= bw
 400         os.close(fd)
 401         return (filepath, timestamp)
 402
 403     def _write_dict_of_arrays(self, d, base_filepath):
 404         """Save dict of (string, numpy_array) pairs under `base_filepath`.
 405
 406         Parameters
 407         ----------
 408         d : dict
 409             Dictionary to save.
 410         base_filepath : path
 411             Path for table of contents and from which per-pair paths
 412             are constructed.
 413         """
 414         # open a new file in readonly mode, don't clobber.
 415         bfd = open(base_filepath, 'w', 0644)
 416         bfd.write("Contents (key : file-extension : format):\n")
 417         for key in d.keys():
 418             clean_key = self._clean_filename(key)
 419             bfd.write("%s : %s : %s\n" % (key, clean_key, str(d[key].dtype)))
 420             # write the keyed array to it's own file
 421             filepath = "%s_%s" % (base_filepath, clean_key)
 422             d[key].tofile(filepath)
 423         bfd.close()
 424
 425     def write_dict_of_arrays(self, d, timestamp=None):
 426         """Save dict of (string, numpy_array) pairs to timestamped files.
 427
 428         Parameters
 429         ----------
 430         d : dict
 431             Dictionary to save.
 432         timestamp : string
 433             Passed on to `._get_filename()`.
 434
 435         Returns
 436         -------
 437         filepath : path
 438             Path to the timestamped log file.
 439         timestamp : string
 440             The timestamp used to generate the log file(s).
 441
 442         Examples
 443         --------
 444         >>> import os
 445         >>> import shutil
 446         >>> import numpy
 447         >>> dl = DataLog(log_dir='test_write_dict_of_arrays',
 448         ...              log_name='my-log', timestamp='20101103')
 449         >>> d = {'data1':numpy.arange(5, dtype=numpy.int16),
 450         ...      'd\/at:$a 2':numpy.arange(3, dtype=numpy.float64)}
 451         >>> filepath,ts = dl.write_dict_of_arrays(
 452         ...     d, timestamp='20101103235959')
 453         >>> filepath
 454         'test_write_dict_of_arrays/20101103/20101103235959_my-log'
 455         >>> print '\\n'.join(sorted(os.listdir(dl._subdir)))
 456         20101103235959_my-log
 457         20101103235959_my-log_data1
 458         20101103235959_my-log_data_2
 459         >>> contents = open(filepath, 'r').read()
 460         >>> print contents
 461         Contents (key : file-extension : format):
 462         data1 : data1 : int16
 463         d\/at:$a 2 : data_2 : float64
 464         <BLANKLINE>
 465         >>> data1_in = numpy.fromfile(
 466         ...     filepath+'_data1', dtype=numpy.int16, count=-1)
 467         >>> data1_in
 468         array([0, 1, 2, 3, 4], dtype=int16)
 469         >>> data2_in = numpy.fromfile(
 470         ...     filepath+'_data_2', dtype=numpy.float64, count=-1)
 471         >>> data2_in
 472         array([ 0.,  1.,  2.])
 473         >>> shutil.rmtree(dl._log_dir)
 474         """
 475         base_filepath,timestamp = self._get_filename(timestamp)
 476         self._write_dict_of_arrays(d, base_filepath)
 477         return (base_filepath, timestamp)
 478
 479
 480 class DataLoad (object):
 481     """Load data logged by `DataLog`.
 482     """
 483     def read(self, filename):
 484         """Load an object saved with `DataLog.write()`.
 485
 486         Parameters
 487         ----------
 488         filename : path
 489             `filename` returned by `DataLog.write()`.
 490
 491         Returns
 492         -------
 493         obj : object
 494             The saved object.
 495
 496         Examples
 497         --------
 498         >>> import shutil
 499         >>> dl = DataLog(log_dir='test_read',
 500         ...              log_name='my-log', timestamp='20101103')
 501         >>> f,t = dl.write([1, 2, 3])
 502         >>> load = DataLoad()
 503         >>> d = load.read(f)
 504         >>> d
 505         [1, 2, 3]
 506         >>> shutil.rmtree(dl._log_dir)
 507         """
 508         return pickle.load(open(filename, 'rb'))
 509
 510     def read_binary(self, filename):
 511         """Load an object saved with `DataLog.write_binary()`.
 512
 513         Warning: this method *requires* `filename` to end with
 514         `_float` and *assumes* that the file contains `numpy.float`
 515         data.  That is terrible.  Use `h5py` instead of this module!
 516
 517         Parameters
 518         ----------
 519         filename : path
 520             `filename` returned by `DataLog.write_binary()`.
 521
 522         Returns
 523         -------
 524         obj : object
 525             The saved object.
 526
 527         Examples
 528         --------
 529         >>> import shutil
 530         >>> import numpy
 531         >>> dl = DataLog(log_dir='test_read_binary',
 532         ...              log_name='my-log_float', timestamp='20101103')
 533         >>> f,t = dl.write_binary(numpy.array([1, 2, 3], dtype=numpy.float))
 534         >>> load = DataLoad()
 535         >>> d = load.read_binary(f)
 536         >>> d
 537         array([ 1.,  2.,  3.])
 538         >>> shutil.rmtree(dl._log_dir)
 539         """
 540         type_ = filename.split("_")[-1]
 541         if type_ == "float":
 542             t = numpy.float
 543         else:
 544             raise Exception(
 545                 "read_binary() not implemented for type %s" % (type_))
 546         return numpy.fromfile(filename, dtype=t)
 547
 548     def read_dict_of_arrays(self, basefile):
 549         """Load an object saved with `DataLog.write_dict_of_arrays()`.
 550
 551         The filenames must not have been altered.
 552
 553         Parameters
 554         ----------
 555         filename : path
 556             `filename` returned by `DataLog.write_dict_of_arrays()`.
 557
 558         Returns
 559         -------
 560         obj : object
 561             The saved object.
 562
 563         Examples
 564         --------
 565         >>> import pprint
 566         >>> import shutil
 567         >>> import numpy
 568         >>> dl = DataLog(log_dir='test_read_dict_of_arrays',
 569         ...              log_name='my-log', timestamp='20101103')
 570         >>> d = {'data1':numpy.arange(5, dtype=numpy.int16),
 571         ...      'd\/at:$a 2':numpy.arange(3, dtype=numpy.float64)}
 572         >>> f,t = dl.write_dict_of_arrays(d, timestamp='20101103235959')
 573         >>> load = DataLoad()
 574         >>> d = load.read_dict_of_arrays(f)
 575         >>> pprint.pprint(d)
 576         {'d\\\\/at:$a 2': array([ 0.,  1.,  2.]),
 577          'data1': array([0, 1, 2, 3, 4], dtype=int16)}
 578         >>> shutil.rmtree(dl._log_dir)
 579         """
 580         obj = {}
 581         i=0
 582         realbasefile = os.path.realpath(basefile)
 583         for line in file(realbasefile):
 584             if i > 0 : # ignore first line
 585                 ldata = line.split(' : ')
 586                 name = ldata[0]
 587                 fpath = "%s_%s" % (realbasefile, ldata[1])
 588                 type_ = getattr(numpy, ldata[2].strip())
 589                 obj[name] = numpy.fromfile(fpath, dtype=type_)
 590             i += 1
 591         return obj
 592
 593
 594 def test():
 595     import doctest
 596     import sys
 597
 598     result = doctest.testmod()
 599     sys.exit(min(result.failed, 127))
 600
 601 if __name__ == "__main__":
 602     test()