stencil_db_from_package with arch_translation_table added

src/copapy/__init__.py

@@ -9,13 +9,21 @@ import platform
 
 Operand = type['Net'] | float | int
 
+
 def get_var_name(var: Any, scope: dict[str, Any] = globals()) -> list[str]:
     return [name for name, value in scope.items() if value is var]
 
-_arch = platform.machine()
-_stencil_data = pkgutil.get_data(__name__, f"obj/stencils_{_arch}_O3.o")
-assert _stencil_data
-generic_sdb = stencil_database(_stencil_data)
+
+def stencil_db_from_package(arch: str = 'native', optimization: str = 'O3') -> stencil_database:
+    arch_translation_table = {'ARM64': 'x86_64'}
+    if arch == 'native':
+        arch = arch_translation_table.get(platform.machine(), platform.machine())
+    stencil_data = pkgutil.get_data(__name__, f"obj/stencils_{arch}_{optimization}.o")
+    assert stencil_data, f"stencils_{arch}_{optimization} not found"
+    return stencil_database(stencil_data)
+
+
+generic_sdb = stencil_db_from_package()
 
 
 class Node:
@@ -414,15 +422,10 @@ def compile_to_instruction_list(node_list: Iterable[Node], sdb: stencil_database
 class Target():
 
     def __init__(self, arch: str = 'native', optimization: str = 'O3') -> None:
-        if arch == 'native':
-            arch = platform.machine()
-        stencil_data = pkgutil.get_data(__name__, f"obj/stencils_{arch}_{optimization}.o")
-        assert stencil_data
-        self.sdb = stencil_database(stencil_data)
+        self.sdb = stencil_db_from_package(arch, optimization)
         self._variables: dict[Net, tuple[int, int, str]] = dict()
 
-
-    def compile(self, *variables: list[Net] | list[list[Net]]) -> None:
+    def compile(self, *variables: Net | list[Net]) -> None:
         nodes: list[Node] = []
         for s in variables:
             if isinstance(s, Net):
@@ -432,12 +435,10 @@ class Target():
                     assert isinstance(net, Net)
                     nodes.append(Write(net))
 
-
         dw, self._variables = compile_to_instruction_list(nodes, self.sdb)
         dw.write_com(binw.Command.END_PROG)
         assert coparun(dw.get_data()) > 0
 
-
     def run(self) -> None:
         # set entry point and run code
         dw = binw.data_writer(self.sdb.byteorder)
@@ -446,7 +447,6 @@ class Target():
         dw.write_com(binw.Command.END_PROG)
         assert coparun(dw.get_data()) > 0
 
-
     def read_value(self, net: Net) -> float | int:
         assert net in self._variables, f"Variable {net} not found"
         addr, lengths, var_type = self._variables[net]

tests/test_coparun_module.py

@@ -1,7 +1,4 @@
-from coparun_module import coparun
-from copapy import Write, const, Target
-import copapy
-from copapy import binwrite
+from copapy import const, Target
 
 
 def test_compile():