src/third_party/pyelftools/elftools/construct/lib/bitstream.py

   1 from .binary import encode_bin, decode_bin
   2
   3 class BitStreamReader(object):
   4
   5     __slots__ = ["substream", "buffer", "total_size"]
   6
   7     def __init__(self, substream):
   8         self.substream = substream
   9         self.total_size = 0
  10         self.buffer = ""
  11
  12     def close(self):
  13         if self.total_size % 8 != 0:
  14             raise ValueError("total size of read data must be a multiple of 8",
  15                 self.total_size)
  16
  17     def tell(self):
  18         return self.substream.tell()
  19
  20     def seek(self, pos, whence = 0):
  21         self.buffer = ""
  22         self.total_size = 0
  23         self.substream.seek(pos, whence)
  24
  25     def read(self, count):
  26         if count < 0:
  27             raise ValueError("count cannot be negative")
  28
  29         l = len(self.buffer)
  30         if count == 0:
  31             data = ""
  32         elif count <= l:
  33             data = self.buffer[:count]
  34             self.buffer = self.buffer[count:]
  35         else:
  36             data = self.buffer
  37             count -= l
  38             bytes = count // 8
  39             if count & 7:
  40                 bytes += 1
  41             buf = encode_bin(self.substream.read(bytes))
  42             data += buf[:count]
  43             self.buffer = buf[count:]
  44         self.total_size += len(data)
  45         return data
  46
  47 class BitStreamWriter(object):
  48
  49     __slots__ = ["substream", "buffer", "pos"]
  50
  51     def __init__(self, substream):
  52         self.substream = substream
  53         self.buffer = []
  54         self.pos = 0
  55
  56     def close(self):
  57         self.flush()
  58
  59     def flush(self):
  60         bytes = decode_bin("".join(self.buffer))
  61         self.substream.write(bytes)
  62         self.buffer = []
  63         self.pos = 0
  64
  65     def tell(self):
  66         return self.substream.tell() + self.pos // 8
  67
  68     def seek(self, pos, whence = 0):
  69         self.flush()
  70         self.substream.seek(pos, whence)
  71
  72     def write(self, data):
  73         if not data:
  74             return
  75         if type(data) is not str:
  76             raise TypeError("data must be a string, not %r" % (type(data),))
  77         self.buffer.append(data)