From 58e831486d3a4e2193a3d1b7ad3c3f4870767871 Mon Sep 17 00:00:00 2001
From: Mohamed Gaber <me@donn.website>
Date: Wed, 19 Nov 2025 03:42:32 +0200
Subject: [PATCH] pyosys: __getitem__ for supported classes

- functions that have a const `[]` operator method now support `__getitem__` in Python
- fields of a pointer type now return a `reference_internal` instead of a `copy` because classes referenced to by pointers typically aren't copyable (e.g. RTLIL::Wire, RTLIL::Module, etc)
- removed duplicate of test_script.py
---
 pyosys/generator.py                           | 11 ++++++--
 tests/arch/ecp5/add_sub.py                    | 20 -------------
 .../{test_script.py => test_ecp5_addsub.py}   |  1 -
 tests/pyosys/test_sigspec_it.py               | 28 +++++++++++++++++++
 4 files changed, 36 insertions(+), 24 deletions(-)
 delete mode 100644 tests/arch/ecp5/add_sub.py
 rename tests/pyosys/{test_script.py => test_ecp5_addsub.py} (99%)
 create mode 100644 tests/pyosys/test_sigspec_it.py

diff --git a/pyosys/generator.py b/pyosys/generator.py
index c0ccdac5a..25f87d570 100644
--- a/pyosys/generator.py
+++ b/pyosys/generator.py
@@ -193,7 +193,7 @@ pyosys_headers = [
             ),
             PyosysClass("SigChunk"),
             PyosysClass("SigBit", hash_expr="s"),
-            PyosysClass("SigSpec", hash_expr="s", denylist={"chunks"}),
+            PyosysClass("SigSpec", hash_expr="s", denylist=frozenset({"chunks"})),
             PyosysClass(
                 "Cell",
                 ref_only=True,
@@ -539,6 +539,8 @@ class PyosysWrapperGenerator(object):
                 python_name_override = "__ne__"
             elif function.operator == "<":
                 python_name_override = "__lt__"
+            elif function.operator == "[]" and function.const:
+                python_name_override = "__getitem__"
             else:
                 return
 
@@ -592,7 +594,10 @@ class PyosysWrapperGenerator(object):
             # care
             return
 
-        has_containers = self.register_containers(field)
+        self.register_containers(field)
+        rvp = "py::return_value_policy::copy"
+        if isinstance(field.type, Pointer):
+            rvp = "py::return_value_policy::reference_internal"
 
         definition_fn = f"def_{'readonly' if field.type.const else 'readwrite'}"
         if field.static:
@@ -604,7 +609,7 @@ class PyosysWrapperGenerator(object):
             f'"{field_python_basename}"',
             f"&{metadata.name}::{field.name}",
         ]
-        def_args.append("py::return_value_policy::copy")
+        def_args.append(rvp)
         print(
             f"\t\t\t.{definition_fn}({', '.join(def_args)})",
             file=self.f,
diff --git a/tests/arch/ecp5/add_sub.py b/tests/arch/ecp5/add_sub.py
deleted file mode 100644
index 0232ac1db..000000000
--- a/tests/arch/ecp5/add_sub.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import os
-from pyosys import libyosys as ys
-
-__dir__ = os.path.dirname(os.path.abspath(__file__))
-add_sub = os.path.join(__dir__, "..", "common", "add_sub.v")
-
-base = ys.Design()
-ys.run_pass(f"read_verilog {add_sub}", base)
-ys.run_pass("hierarchy -top top", base)
-ys.run_pass("proc", base)
-ys.run_pass("equiv_opt -assert -map +/ecp5/cells_sim.v synth_ecp5", base)
-
-postopt = ys.Design()
-ys.run_pass("design -load postopt", postopt)
-ys.run_pass("cd top", postopt)
-ys.run_pass("select -assert-min 25 t:LUT4", postopt)
-ys.run_pass("select -assert-max 26 t:LUT4", postopt)
-ys.run_pass("select -assert-count 10 t:PFUMX", postopt)
-ys.run_pass("select -assert-count 6 t:L6MUX21", postopt)
-ys.run_pass("select -assert-none t:LUT4 t:PFUMX t:L6MUX21 %% t:* %D", postopt)
diff --git a/tests/pyosys/test_script.py b/tests/pyosys/test_ecp5_addsub.py
similarity index 99%
rename from tests/pyosys/test_script.py
rename to tests/pyosys/test_ecp5_addsub.py
index 7c3ec96ef..ddc50b775 100644
--- a/tests/pyosys/test_script.py
+++ b/tests/pyosys/test_ecp5_addsub.py
@@ -1,7 +1,6 @@
 from pathlib import Path
 from pyosys import libyosys as ys
 
-
 __file_dir__ = Path(__file__).absolute().parent
 add_sub = __file_dir__.parent / "arch" / "common" / "add_sub.v"
 
diff --git a/tests/pyosys/test_sigspec_it.py b/tests/pyosys/test_sigspec_it.py
new file mode 100644
index 000000000..2876e7725
--- /dev/null
+++ b/tests/pyosys/test_sigspec_it.py
@@ -0,0 +1,28 @@
+from pyosys import libyosys as ys
+from pathlib import Path
+
+__file_dir__ = Path(__file__).absolute().parent
+
+def _dump_sigbit(bit):
+    if bit.is_wire():
+        if bit.wire.width == 1:
+            return bit.wire.name.str()
+        else:
+            return f"{bit.wire.name} [{bit.offset}]"
+    else:
+        if bit.data == ys.State.S1:
+            return 1
+        elif bit.data == ys.State.S0:
+            return 0
+        else:
+            assert "unknown constants not supported"
+
+d = ys.Design()
+
+ys.run_pass(f"read_verilog {__file_dir__ / 'spm.cut.v.gz'}", d)
+ys.run_pass(f"hierarchy -top spm", d)
+module = d.module(r"\spm")
+for conn_from, conn_to in module.connections_:
+    for bit_from, bit_to in zip(conn_from, conn_to):
+        print(f"assign {_dump_sigbit(bit_from)} = {_dump_sigbit(bit_to)};")
+