From 6df616e40ed871ea55e4271e38872743077e712d Mon Sep 17 00:00:00 2001
From: awygle <awygle@gmail.com>
Date: Mon, 6 Jul 2020 22:17:03 -0700
Subject: [PATCH] hdl.ast: don't inherit Shape from NamedTuple.

Fixes #421.
---
 nmigen/hdl/ast.py           | 34 ++++++++++++++++++++++++----------
 nmigen/test/test_hdl_ast.py | 10 ++++++++++
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/nmigen/hdl/ast.py b/nmigen/hdl/ast.py
index 37a3b46..55b825b 100644
--- a/nmigen/hdl/ast.py
+++ b/nmigen/hdl/ast.py
@@ -32,7 +32,7 @@ class DUID:
         DUID.__next_uid += 1
 
 
-class Shape(typing.NamedTuple):
+class Shape:
     """Bit width and signedness of a value.
 
     A ``Shape`` can be constructed using:
@@ -55,8 +55,15 @@ class Shape(typing.NamedTuple):
     signed : bool
         If ``False``, the value is unsigned. If ``True``, the value is signed two's complement.
     """
-    width:  int  = 1
-    signed: bool = False
+    def __init__(self, width=1, signed=False):
+        if not isinstance(width, int) or width < 0:
+            raise TypeError("Width must be a non-negative integer, not {!r}"
+                            .format(width))
+        self.width = width
+        self.signed = signed
+
+    def __iter__(self):
+        return iter((self.width, self.signed))
 
     @staticmethod
     def cast(obj, *, src_loc_at=0):
@@ -95,13 +102,20 @@ class Shape(typing.NamedTuple):
         else:
             return "unsigned({})".format(self.width)
 
-
-# TODO: use dataclasses instead of this hack
-def _Shape___init__(self, width=1, signed=False):
-    if not isinstance(width, int) or width < 0:
-        raise TypeError("Width must be a non-negative integer, not {!r}"
-                        .format(width))
-Shape.__init__ = _Shape___init__
+    def __eq__(self, other):
+        if isinstance(other, tuple) and len(other) == 2:
+            width, signed = other
+            if isinstance(width, int) and isinstance(signed, bool):
+                return self.width == width and self.signed == signed
+            else:
+                raise TypeError("Shapes may be compared with other Shapes and (int, bool) tuples, "
+                        "not {!r}"
+                        .format(other))
+        if not isinstance(other, Shape):
+            raise TypeError("Shapes may be compared with other Shapes and (int, bool) tuples, "
+                    "not {!r}"
+                    .format(other))
+        return self.width == other.width and self.signed == other.signed
 
 
 def unsigned(width):
diff --git a/nmigen/test/test_hdl_ast.py b/nmigen/test/test_hdl_ast.py
index b724ad6..691408a 100644
--- a/nmigen/test/test_hdl_ast.py
+++ b/nmigen/test/test_hdl_ast.py
@@ -39,6 +39,16 @@ class ShapeTestCase(FHDLTestCase):
                 msg="Width must be a non-negative integer, not -1"):
             Shape(-1)
 
+    def test_compare_wrong(self):
+        with self.assertRaises(TypeError,
+                msg="Shapes may be compared with other Shapes and (int, bool) tuples, not 'hi'"):
+            Shape(1, True) == 'hi'
+
+    def test_compare_tuple_wrong(self):
+        with self.assertRaises(TypeError,
+                msg="Shapes may be compared with other Shapes and (int, bool) tuples, not (2, 3)"):
+            Shape(1, True) == (2, 3)
+
     def test_repr(self):
         self.assertEqual(repr(Shape()), "unsigned(1)")
         self.assertEqual(repr(Shape(2, True)), "signed(2)")
-- 
2.30.2