Add cases for Bits, UInt, SInt

phanrahan · May 23, 2019 · f14f82b · f14f82b
1 parent 52c43e3
commit f14f82b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 8 deletions.
diff --git a/magma/util.py b/magma/util.py
@@ -24,8 +24,13 @@ def pretty_str(t):
         # Insert first newline + indent and last newline
         result = "\n    " + result + "\n"
         s = f"Tuple({result})"
+    elif isinstance(t, m.SIntKind):
+        s = f"SInt[{t.N}]"
+    elif isinstance(t, m.UIntKind):
+        s = f"UInt[{t.N}]"
+    elif isinstance(t, m.BitsKind):
+        s = f"Bits[{t.N}]"
     elif isinstance(t, m.ArrayKind):
-        s = "Array[%d, %s]" % (t.N, t.T)
         s = f"Array[{t.N}, {pretty_str(t.T)}]"
     else:
         s = str(t)

diff --git a/tests/test_type/test_pretty_print.py b/tests/test_type/test_pretty_print.py
@@ -44,20 +44,20 @@ def test_pretty_print_array_of_tuple():
 
 
 def test_pretty_print_array_of_nested_tuple():
-    t = m.Tuple(a=m.Bit, b=m.Bit, c=m.Bit)
+    t = m.Tuple(a=m.Bits[5], b=m.UInt[3], c=m.SInt[4])
     u = m.Tuple(x=t, y=t)
     v = m.Array[3, u]
     assert m.util.pretty_str(v) == """\
 Array[3, Tuple(
     x = Tuple(
-        a = Bit,
-        b = Bit,
-        c = Bit
+        a = Bits[5],
+        b = UInt[3],
+        c = SInt[4]
     ),
     y = Tuple(
-        a = Bit,
-        b = Bit,
-        c = Bit
+        a = Bits[5],
+        b = UInt[3],
+        c = SInt[4]
     )
 )]\
 """