records.py

   1 import struct, os, re
   2
   3 from logging import debug
   4
   5 BIG_ENDIAN = True
   6 LITTLE_ENDIAN = False
   7
   8 class RecordsException(Exception):
   9     pass
  10
  11 class FormatError(RecordsException):
  12     pass
  13
  14 class UnknownSizeException(RecordsException):
  15     pass
  16
  17 ### Fields
  18 class BaseField:
  19     def set_bigendian(self, ignore): pass
  20     def set_reclen(self, ignore): pass
  21     def put_default(self, dict): pass
  22
  23 ###### Skip is a dummy field that isn't read
  24 class Skip(BaseField):
  25     def __init__(self, n):
  26         self.get_size = (lambda:n)
  27         self.write = self.read
  28     def read(self, file, dict):
  29         file.seek(self.get_size(), os.SEEK_CUR)
  30
  31 ###### Field is a composite field object factory
  32 class Field(BaseField):
  33     # Value handlers - know what to do with read values and
  34     # where to get values to be written
  35     class Named:
  36         def __init__(self, name, default):
  37             self.name = name
  38             self.default = default
  39         def get(self, dict):
  40             return dict[self.name] if dict.has_key(self.name) else self.default
  41         def put(self, dict, value):
  42             dict[self.name] = value
  43         def put_default(self, dict):
  44             dict[self.name] = self.default
  45         def __str__(self): return self.name
  46
  47     class LenLookup:
  48         re = re.compile("len\(([^)]+)\)")
  49         def __init__(self, expr):
  50             self.expr = expr
  51             self.array_name = self.re.match(expr).group(1)
  52         def get(self, dict):
  53             return len( dict[self.array_name] )
  54         def put(self, dict, value):
  55             dict[self.expr] = value
  56         def put_default(self, dict): pass
  57         def __str__(self): return self.expr
  58
  59     class Const:
  60         def __init__(self, const, check):
  61             self.const = const
  62             if check:
  63                 self.put = self.check
  64         def get(self, dict):
  65             return self.const
  66         def check(self, dict, value):
  67             if value != self.const:
  68                 raise FormatError("Expected %s, got %s" % (self.const, value))
  69         def put(self, dict, value): pass
  70         def put_default(self, dict): pass
  71         def __str__(self): return "Constant(%s)" % self.const
  72
  73     # Field factory - composes field from packer and value handler
  74     def __init__(self, packer, name=None, default=None, const=None, check = False):
  75         # helpers
  76         def set_value_handler(vh):
  77             self.put = vh.put
  78             self.get = vh.get
  79             self.put_default = vh.put_default
  80             self.__str__ = vh.__str__
  81         def const_later(const):
  82             set_value_handler(Field.Const(const, check))
  83
  84         if name == '%reclen%':
  85             self.set_reclen = const_later
  86         elif name is not None:
  87             if Field.LenLookup.re.match(name) is not None:
  88                 set_value_handler(Field.LenLookup(name))
  89             else:
  90                 set_value_handler(Field.Named(name,default))
  91         elif const is not None:
  92             set_value_handler(Field.Const(const, check))
  93         else:
  94             raise RecordsException("Bad field parameters")
  95
  96         self.get_size = packer.get_size
  97         self.pack = packer.pack
  98         self.unpack = packer.unpack
  99         if packer.__class__.__dict__.has_key('set_bigendian'):
 100             self.set_bigendian = packer.set_bigendian
 101
 102     def read(self, file, dict):
 103         value = self.unpack( file.read( self.get_size() ))
 104         debug("read %s: %s", self, value)
 105         self.put(dict, value)
 106     def write(self, file, dict):
 107         value = self.get(dict)
 108         debug("write %s: %s", self, value)
 109         file.write( self.pack( value ))
 110
 111 class ListField(BaseField):
 112     def __init__(self, name, cell):
 113         self.cell = cell
 114         self.name = name
 115
 116     def read(self, file, dict):
 117         lst = dict[self.name] if dict.has_key(self.name) else []
 118         n = dict["len(%s)" % self.name]
 119         for i in xrange(0, n):
 120             if i < len(lst):
 121                 lst[i] = self.cell.read(file, lst[i])
 122             else:
 123                 lst.append(self.cell.read(file, {}))
 124         dict[self.name] = lst
 125
 126     def write(self, file, dict):
 127         lst = dict[ self.name ]
 128         for c in lst:
 129             self.cell.write( file, c )
 130
 131     def get_size(self):
 132         raise UnknownSizeException()
 133
 134     def set_bigendian(self, bigendian):
 135         self.cell.set_bigendian(bigendian)
 136
 137 ### Packers
 138 class SimplePacker:
 139     def __init__(self, fmt):
 140         self.fmt = fmt
 141         self.size = struct.calcsize(fmt)
 142     def pack(self,val): return struct.pack(self.fmt,val)
 143     def unpack(self,str): return struct.unpack(self.fmt,str)[0]
 144     def get_size(self): return self.size
 145
 146 class Uint8(SimplePacker):
 147     def __init__(self): SimplePacker.__init__(self,"B")
 148
 149 class Bool8(Uint8):
 150     def pack(self, val): return Uint8.pack(self, (1 if val else 0))
 151     def unpack(self, str): return Uint8.unpack(self, str) != 0
 152
 153 class Uint24:
 154     def __init__(self, bigendian = LITTLE_ENDIAN):
 155         self.bigendian = bigendian
 156     def get_size(self): return 3
 157     def pack(self, i):
 158         if self.bigendian:
 159             return struct.pack(">I",i)[1:4]
 160         else:
 161             return struct.pack("<I",i)[0:3]
 162     def unpack(self, s):
 163         if self.bigendian:
 164              return struct.unpack('>I','\x00' + s[0:3])[0]
 165         else:
 166             return struct.unpack('<I',s[0:3] + '\x00')[0]
 167     def set_bigendian(self, bigendian):
 168         self.bigendian = bigendian
 169
 170 class Int24(Uint24):
 171     def __init__(self, bigendian = LITTLE_ENDIAN):
 172         self.bigendian = bigendian
 173     def pack(self, i):
 174         if self.bigendian:
 175             return struct.pack(">i",i)[1:4]
 176         else:
 177             return struct.pack("<i",i)[0:3]
 178     def unpack(self, s):
 179         u = Uint24.unpack(self,s)
 180         if (u & 0x800) != 0:
 181             return - ((~u + 1) & 0xfff)
 182         else:
 183             return u
 184
 185 class Bool24(Int24):
 186     def __init__(self): Int24.__init__(self)
 187     def pack(self, val): return Int24.pack(self, (-1 if val else 0))
 188     def unpack(self, str): return Int24.unpack(self, str) != 0
 189
 190 class ZeroPaddedString:
 191     def __init__(self, len, enc):
 192         self.size = len
 193         self.enc = enc
 194     def pack(self, val):
 195         return val.encode(self.enc).ljust(self.size,'\x00')
 196     def unpack(self, str):
 197         if len(str) != self.size:
 198             raise FormatError()
 199         return str.decode(self.enc).rstrip('\x00')
 200     def get_size(self): return self.size
 201
 202 ### Record - an ordered list of fields
 203 class Record(BaseField):
 204     def __init__(self, fields, bigendian=None):
 205         self.fields = fields
 206         try:
 207             reclen = self.get_size()
 208             for f in fields:
 209                 f.set_reclen(reclen)
 210         except UnknownSizeException:
 211             pass
 212         if bigendian is not None:
 213             self.set_bigendian(bigendian)
 214
 215     def read(self, file, dict=None):
 216         if dict is None: dict = {}
 217         for f in self.fields:
 218             f.read(file, dict)
 219         return dict
 220
 221     def write(self, file, dict):
 222         for f in self.fields:
 223             f.write(file, dict)
 224
 225     def make_default(self, dict = None):
 226         if dict is None:
 227             dict = {}
 228         for f in self.fields:
 229             f.put_default(dict)
 230         return dict
 231
 232     def get_size(self):
 233         size = 0
 234         for f in self.fields:
 235             size += f.get_size()
 236         return size
 237
 238     def set_bigendian(self, bigendian):
 239         for f in self.fields:
 240             f.set_bigendian(bigendian)
 241