ebml/schema/base.py

   1 import abc
   2 try:
   3         from cStringIO import StringIO
   4 except ImportError:
   5         from StringIO import StringIO
   6 from ..core import *
   7
   8
   9 __all__ = ('UnknownElement', 'Element', 'Document', 'INT', 'UINT', 'FLOAT', 'STRING', 'UNICODE', 'DATE', 'BINARY', 'CONTAINER')
  10
  11
  12 INT, UINT, FLOAT, STRING, UNICODE, DATE, BINARY, CONTAINER = range(0, 8)
  13
  14
  15 READERS = {
  16         INT: read_signed_integer,
  17         UINT: read_unsigned_integer,
  18         FLOAT: read_float,
  19         STRING: read_string,
  20         UNICODE: read_unicode_string,
  21         DATE: read_date,
  22         BINARY: lambda stream, size: bytearray(stream.read(size))
  23 }
  24
  25
  26 ENCODERS = {
  27         INT: encode_signed_integer,
  28         UINT: encode_unsigned_integer,
  29         FLOAT: encode_float,
  30         STRING: encode_string,
  31         UNICODE: encode_unicode_string,
  32         DATE: encode_date,
  33         BINARY: lambda binary, length: binary
  34 }
  35
  36
  37 VALIDATORS = {
  38         INT: lambda value: True if isinstance(value, (int, long)) else False,
  39         UINT: lambda value: True if isinstance(value, (int, long)) and value == abs(value) else False,
  40         FLOAT: lambda value: True if isinstance(value, float) else False,
  41         STRING: lambda value: True if isinstance(value, str) else False,
  42         UNICODE: lambda value: True if isinstance(value, basestring) else False,
  43         DATE: lambda value: True if isinstance(value, datetime.datetime) else False,
  44         BINARY: lambda value: True if isinstance(value, (str, bytes, bytearray)) else False
  45 }
  46
  47
  48 class BaseElement(object):
  49         __metaclass__ = abc.ABCMeta
  50
  51         id = abc.abstractproperty()
  52         name = abc.abstractproperty()
  53         type = abc.abstractproperty()
  54         default = None
  55         children = ()
  56         mandatory = False
  57         multiple = False
  58
  59
  60 class UnknownElement(BaseElement):
  61         id = None
  62         name = 'Unknown'
  63         type = BINARY
  64
  65         def __init__(self, id, encoding):
  66                 self.id = id
  67                 self.encoding = encoding
  68
  69
  70 def read_elements(stream, size, document, children):
  71         elements = []
  72         while (size if size is not None else True):
  73                 try:
  74                         element_id, element_id_size = read_element_id(stream)
  75                         element_size, element_size_size = read_element_size(stream)
  76                         element_encoding = (element_size, bytearray(stream.read(element_size)))
  77                 except:
  78                         break
  79                 else:
  80                         element_class = None
  81                         for child in (children + document.globals):
  82                                 if child.id == element_id:
  83                                         element_class = child
  84                                         break
  85                         if element_class is None:
  86                                 element = UnknownElement(element_id, element_encoding)
  87                         else:
  88                                 element = element_class(document, encoding=element_encoding)
  89                         elements.append(element)
  90                         if size is not None:
  91                                 size -= element_id_size + element_size_size + element_size
  92         return elements
  93
  94
  95 class Element(BaseElement):
  96         @classmethod
  97         def check_value(cls, value):
  98                 if cls.type in VALIDATORS:
  99                         return VALIDATORS[cls.type](value)
 100                 elif cls.type == CONTAINER:
 101                         if isinstance(value, (list, tuple)):
 102                                 for item in value:
 103                                         if not isinstance(value, Element):
 104                                                 return False
 105                                 return True
 106                         elif isinstance(value, Element):
 107                                 return True
 108                         else:
 109                                 return False
 110                 else:
 111                         raise NotImplementedError('Unsupported element type.')
 112
 113         def __init__(self, document, value=None, encoding=None):
 114                 self.document = document
 115                 self._value = value
 116                 self._encoding = encoding
 117
 118         @property
 119         def value(self):
 120                 if self._value is None and self._encoding is not None:
 121                         if self.type in READERS:
 122                                 self._value = READERS[self.type](StringIO(self._encoding[1]), self._encoding[0])
 123                         elif self.type == CONTAINER:
 124                                 self._value = read_elements(StringIO(self._encoding[1]), self._encoding[0], self.document, self.children)
 125                 return self._value
 126
 127         @value.setter
 128         def set_value(self, value):
 129                 if not self.check_value(value):
 130                         raise ValueError('Unsupported element value.')
 131                 self._value = value
 132                 self._encoding = None
 133
 134         @property
 135         def encoding(self):
 136                 if self._encoding is None:
 137                         size = 0
 138                         data = bytearray()
 139                         if self._value is not None:
 140                                 if self.type in ENCODERS:
 141                                         data = ENCODERS[self.type](self._value)
 142                                         size = len(data)
 143                                 elif self.type == CONTAINER:
 144                                         for element in self._value:
 145                                                 size += element.size
 146                                                 data.extend(element.encoding[1])
 147                         self._encoding = (size, data)
 148                 return self._encoding
 149
 150         @property
 151         def id_size(self):
 152                 return len(encode_element_id(self.id))
 153
 154         @property
 155         def size_size(self):
 156                 return len(encode_element_size(self.body_size))
 157
 158         @property
 159         def head_size(self):
 160                 return self.id_size + self.size_size
 161
 162         @property
 163         def body_size(self):
 164                 return self.encoding[0]
 165
 166         @property
 167         def size(self):
 168                 return self.head_size + self.body_size
 169
 170
 171 class Document(object):
 172         __metaclass__ = abc.ABCMeta
 173
 174         type = abc.abstractproperty()
 175         version = abc.abstractproperty()
 176         children = ()
 177         globals = ()
 178
 179         def __init__(self, stream):
 180                 self.stream = stream
 181                 self._roots = None
 182
 183         @property
 184         def roots(self):
 185                 if self._roots is None:
 186                         self._roots = read_elements(self.stream, None, self, self.children)
 187                 return self._roots