src/soc/decoder/selectable_int.py

   1 import unittest
   2 class SelectableInt:
   3     def __init__(self, value, bits):
   4         self.value = value
   5         self.bits = bits
   6
   7     def __add__(self, b):
   8         assert b.bits == self.bits
   9         return SelectableInt(self.value + b.value, self.bits)
  10
  11     def __getitem__(self, key):
  12         if isinstance(key, int):
  13             assert key < self.bits
  14             assert key >= 0
  15
  16             value = (self.value >> key) & 1
  17             return SelectableInt(value, 1)
  18         elif isinstance(key, slice):
  19             assert key.step is None or key.step == 1
  20             assert key.start < key.stop
  21             assert key.start >= 0
  22             assert key.stop <= self.bits
  23
  24             bits = key.stop - key.start
  25             mask = (1 << bits) - 1
  26             value = (self.value >> key.start) & mask
  27             return SelectableInt(value, bits)
  28
  29     def __eq__(self, other):
  30         if isinstance(other, SelectableInt):
  31             return other.value == self.value and other.bits == self.bits
  32         if isinstance(other, int):
  33             return other == self.value
  34         assert False
  35
  36     def __repr__(self):
  37         return "SelectableInt(value={:x}, bits={})".format(self.value, self.bits)
  38
  39
  40 class SelectableIntTestCase(unittest.TestCase):
  41     def test_add(self):
  42         a = SelectableInt(5, 8)
  43         b = SelectableInt(9, 8)
  44         c = a + b
  45         assert c.value == a.value + b.value
  46         assert c.bits == a.bits
  47
  48     def test_select(self):
  49         a = SelectableInt(0xa5, 8)
  50         assert a[0] == 1
  51         assert a[0:1] == 1
  52         assert a[0:4] == 5
  53         assert a[4:8] == 10
  54
  55
  56 if __name__ == "__main__":
  57     unittest.main()