From: whitequark <cz@m-labs.hk>
Date: Fri, 19 Apr 2019 19:55:39 +0000 (+0000)
Subject: hdl.rec: fix slicing of records.
X-Git-Url: https://git.libre-soc.org/?a=commitdiff_plain;h=c4810e6da07e42db9575b221b11e6172f6a04946;p=nmigen.git

hdl.rec: fix slicing of records.
---

diff --git a/nmigen/hdl/rec.py b/nmigen/hdl/rec.py
index f1ddbd2..8f1065b 100644
--- a/nmigen/hdl/rec.py
+++ b/nmigen/hdl/rec.py
@@ -85,16 +85,19 @@ class Record(Value):
     def __getattr__(self, name):
         return self[name]
 
-    def __getitem__(self, name):
-        try:
-            return self.fields[name]
-        except KeyError:
-            if self.name is None:
-                reference = "Unnamed record"
-            else:
-                reference = "Record '{}'".format(self.name)
-            raise NameError("{} does not have a field '{}'. Did you mean one of: {}?"
-                            .format(reference, name, ", ".join(self.fields))) from None
+    def __getitem__(self, item):
+        if isinstance(item, str):
+            try:
+                return self.fields[item]
+            except KeyError:
+                if self.name is None:
+                    reference = "Unnamed record"
+                else:
+                    reference = "Record '{}'".format(self.name)
+                raise NameError("{} does not have a field '{}'. Did you mean one of: {}?"
+                                .format(reference, item, ", ".join(self.fields))) from None
+        else:
+            return super().__getitem__(item)
 
     def shape(self):
         return sum(len(f) for f in self.fields.values()), False
diff --git a/nmigen/test/test_hdl_rec.py b/nmigen/test/test_hdl_rec.py
index 65e8bf6..847eedc 100644
--- a/nmigen/test/test_hdl_rec.py
+++ b/nmigen/test/test_hdl_rec.py
@@ -79,6 +79,15 @@ class RecordTestCase(FHDLTestCase):
         self.assertEqual(repr(r), "(rec <unnamed> stb)")
         self.assertEqual(r.stb.name, "stb")
 
+    def test_iter(self):
+        r = Record([
+            ("data", 4),
+            ("stb",  1),
+        ])
+
+        self.assertEqual(repr(r[0]),   "(slice (rec r data stb) 0:1)")
+        self.assertEqual(repr(r[0:3]), "(slice (rec r data stb) 0:3)")
+
     def test_wrong_field(self):
         r = Record([
             ("stb", 1),