From 586dd42dffcafbdf9743de5404cbbf2b6e159e61 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Tue, 2 Jul 2019 17:44:55 +0000
Subject: [PATCH] hdl.rec: implement slicing by component names.

Fixes #121.
---
 nmigen/hdl/rec.py           | 17 +++++++++++++++--
 nmigen/test/test_hdl_rec.py | 19 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/nmigen/hdl/rec.py b/nmigen/hdl/rec.py
index 2c6694c..67c77c2 100644
--- a/nmigen/hdl/rec.py
+++ b/nmigen/hdl/rec.py
@@ -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)
 
diff --git a/nmigen/test/test_hdl_rec.py b/nmigen/test/test_hdl_rec.py
index 9c72a16..7821c7d 100644
--- a/nmigen/test/test_hdl_rec.py
+++ b/nmigen/test/test_hdl_rec.py
@@ -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):
-- 
2.30.2