hdl.rec: implement slicing by component names.

Fixes #121.
This commit is contained in:
whitequark 2019-07-02 17:44:55 +00:00
parent 34f110100a
commit 6b843b5be6
2 changed files with 34 additions and 2 deletions

View file

@ -56,8 +56,15 @@ class Layout:
shape = (shape, False)
self.fields[name] = (shape, direction)
def __getitem__(self, name):
return self.fields[name]
def __getitem__(self, item):
if isinstance(item, tuple):
return Layout([
(name, shape, dir)
for (name, (shape, dir)) in self.fields.items()
if name in item
])
return self.fields[item]
def __iter__(self):
for name, (shape, dir) in self.fields.items():
@ -121,6 +128,12 @@ class Record(Value):
reference = "Record '{}'".format(self.name)
raise AttributeError("{} does not have a field '{}'. Did you mean one of: {}?"
.format(reference, item, ", ".join(self.fields))) from None
elif isinstance(item, tuple):
return Record(self.layout[item], fields={
field_name: field_value
for field_name, field_value in self.fields.items()
if field_name in item
})
else:
return super().__getitem__(item)

View file

@ -25,6 +25,18 @@ class LayoutTestCase(FHDLTestCase):
self.assertEqual(sublayout["a"], ((1, False), DIR_NONE))
self.assertEqual(sublayout["b"], ((1, False), DIR_NONE))
def test_slice_tuple(self):
layout = Layout.wrap([
("a", 1),
("b", 2),
("c", 3)
])
expect = Layout.wrap([
("a", 1),
("c", 3)
])
self.assertEqual(layout["a", "c"], expect)
def test_wrong_field(self):
with self.assertRaises(TypeError,
msg="Field (1,) has invalid layout: should be either (name, shape) or "
@ -139,6 +151,13 @@ class RecordTestCase(FHDLTestCase):
r4 = Record.like(r1, name_suffix="foo")
self.assertEqual(r4.name, "r1foo")
def test_slice_tuple(self):
r1 = Record([("a", 1), ("b", 2), ("c", 3)])
r2 = r1["a", "c"]
self.assertEqual(r2.layout, Layout([("a", 1), ("c", 3)]))
self.assertIs(r2.a, r1.a)
self.assertIs(r2.c, r1.c)
class ConnectTestCase(FHDLTestCase):
def setUp_flat(self):