Source Code for Module mvpa.misc.io.hamster

  1  # emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- 
  2  # vi: set ft=python sts=4 ts=4 sw=4 et: 
  3  ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## 
  4  # 
  5  #   See COPYING file distributed along with the PyMVPA package for the 
  6  #   copyright and license terms. 
  7  # 
  8  ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ## 
  9  """Helper for simple storage facility via cPickle and optionally zlib""" 
 10   
 11  __docformat__ = 'restructuredtext' 
 12   
 13  import os 
 14   
 15  from mvpa.base import externals 
 16   
 17  _d_geti_ = dict.__getitem__ 
 18  _d_seti_ = dict.__setitem__ 
 19   
 20  _o_geta_ = dict.__getattribute__ 
 21  _o_seta_ = dict.__setattr__ 
 22   
 23  if externals.exists('cPickle', raiseException=True) and \ 
 24     externals.exists('gzip', raiseException=True): 
 25      import cPickle, gzip 
 26   
 27  if __debug__: 
 28      from mvpa.base import debug 
 29   
 30 -class Hamster(object): 
 31      """Simple container class with basic IO capabilities. 
 32   
 33      It is capable of storing itself in a file, or loading from a file using 
 34      cPickle (optionally via zlib from compressed files). Any serializable 
 35      object can be bound to a hamster to be stored. 
 36   
 37      To undig burried hamster use Hamster(filename). Here is an example: 
 38   
 39        >>> h = Hamster(bla='blai') 
 40        >>> h.boo = N.arange(5) 
 41        >>> h.dump(filename) 
 42        ... 
 43        >>> h = Hamster(filename) 
 44   
 45      Since Hamster introduces methods `dump`, `asdict` and property 
 46      'registered', those names cannot be used to assign an attribute, 
 47      nor provided in among constructor arguments. 
 48      """ 
 49   
 50      __ro_attr = set(object.__dict__.keys() + 
 51                      ['dump', 'registered', 'asdict']) 
 52      """Attributes which come with being an object""" 
 53   
 54 -    def __new__(cls, *args, **kwargs): 
 55          if len(args) > 0: 
 56              if len(kwargs) > 0: 
 57                  raise ValueError, \ 
 58                        "Do not mix positional and keyword arguments. " \ 
 59                        "Use a single positional argument -- filename, " \ 
 60                        "or any number of keyword arguments, without having " \ 
 61                        "filename specified" 
 62              if len(args) == 1 and isinstance(args[0], basestring): 
 63                  filename = args[0] 
 64                  args = args[1:] 
 65                  if __debug__: 
 66                      debug('IOH', 'Undigging hamster from %s' % filename) 
 67                  # compressed or not -- that is the question 
 68                  if filename.endswith('.gz'): 
 69                      f = gzip.open(filename) 
 70                  else: 
 71                      f = open(filename) 
 72                  result = cPickle.load(f) 
 73                  if not isinstance(result, Hamster): 
 74                      warning("Loaded other than Hamster class from %s" % filename) 
 75                  return result 
 76              else: 
 77                  raise ValueError, "Hamster accepts only a single positional " \ 
 78                        "argument and it must be a filename. Got %d " \ 
 79                        "arguments" % (len(args),) 
 80          else: 
 81              return object.__new__(cls) 
 82   
 83   
 84 -    def __init__(self, *args, **kwargs): 
 85          """Initialize Hamster. 
 86   
 87          Providing a single parameter string would treat it as a 
 88          filename from which to undig the data. Otherwise all keyword 
 89          parameters are assigned into the attributes of the object. 
 90          """ 
 91          if len(args) > 0: 
 92              if len(args) == 1 and isinstance(args[0], basestring): 
 93                  # it was a filename 
 94                  args = args[1:] 
 95              else: 
 96                  raise RuntimeError, "Should not get here" 
 97   
 98          # assign provided attributes 
 99          for k,v in kwargs.iteritems(): 
100              setattr(self, k, v) 
101   
102          object.__init__(self) 
103   
104   
105 -    def dump(self, filename, compresslevel='auto'): 
106          """Bury the hamster into the file 
107   
108          :Parameter: 
109            filename: str 
110              Name of the target file. When writing to a compressed file the 
111              filename gets a '.gz' extension if not already specified. This 
112              is necessary as the constructor uses the extension to decide 
113              whether it loads from a compressed or uncompressed file. 
114            compresslevel: 'auto' or int 
115              Compression level setting passed to gzip. When set to 
116              'auto', if filename ends with '.gz' `compresslevel` is set 
117              to 5, 0 otherwise.  However, when `compresslevel` is set to 
118              0 gzip is bypassed completely and everything is written to 
119              an uncompressed file. 
120          """ 
121          if compresslevel == 'auto': 
122              compresslevel = (0, 5)[int(filename.endswith('.gz'))] 
123          if compresslevel > 0 and not filename.endswith('.gz'): 
124              filename += '.gz' 
125          if __debug__: 
126              debug('IOH', 'Burying hamster into %s' % filename) 
127          if compresslevel == 0: 
128              f = open(filename, 'w') 
129          else: 
130              f = gzip.open(filename, 'w', compresslevel) 
131          cPickle.dump(self, f) 
132          f.close() 
133   
134   
135 -    def __repr__(self): 
136          reg_attr = ["%s=%s" % (k, repr(getattr(self, k))) 
137                      for k in self.registered] 
138          return "%s(%s)" % (self.__class__.__name__, 
139                             ", ".join(reg_attr)) 
140   
141      # ??? actually seems to be ugly 
142      #def __str__(self): 
143      #    registered = self.registered 
144      #    return "%s with %d elements: %s" \ 
145      #           % (self.__class__.__name__, 
146      #              len(registered), 
147      #              ", ".join(self.registered)) 
148   
149      @property 
150 -    def registered(self): 
151          """List registered attributes. 
152          """ 
153          reg_attr = [k for k in self.__dict__.iterkeys() 
154                      if not k in self.__ro_attr] 
155          reg_attr.sort() 
156          return reg_attr 
157   
158   
159 -    def __setattr__(self, k, v): 
160          """Just to prevent resetting read-only attributes, such as methods 
161          """ 
162          if k in self.__ro_attr: 
163              raise ValueError, "'%s' object attribute '%s' is read-only" \ 
164                    % (self.__class__.__name__, k) 
165          object.__setattr__(self, k, v) 
166   
167   
168 -    def asdict(self): 
169          """Return registered data as dictionary 
170          """ 
171          return dict([(k, getattr(self, k)) 
172                       for k in self.registered]) 
173