From 6a3da2abdbc4d06f59838ab42a5a9c1fbc28cf0d Mon Sep 17 00:00:00 2001
From: Nikolay Repin <nikolat-repin201303@mail.ru>
Date: Tue, 27 Aug 2024 17:31:27 +0700
Subject: [PATCH] [cocas] Implement binary output format for images

---
 cocas/linker/image.py | 41 +++++++++++++++++++++++++++++++++++------
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/cocas/linker/image.py b/cocas/linker/image.py
index 1aecd333..d247bd8d 100644
--- a/cocas/linker/image.py
+++ b/cocas/linker/image.py
@@ -1,14 +1,16 @@
+from collections.abc import Callable
 from pathlib import Path
 from typing import Union
 
+WriterFn = Callable[[Path, bytearray], None]
 
-def write_image(filename: Union[Path, str], arr: bytearray):
-    """
-    Write the contents or array into file in logisim-compatible format
 
-    :param filename: path to output file
-    :param arr: bytearray to be written
-    """
+def write_binary_image(filename: Path, arr: bytearray):
+    with open(filename, mode='wb') as f:
+        f.write(arr)
+
+
+def write_logisim_image(filename: Path, arr: bytearray):
     with open(filename, mode='w') as f:
         f.write("v2.0 raw\n")
         zeroes = 0
@@ -25,3 +27,30 @@ def write_image(filename: Union[Path, str], arr: bytearray):
                             f.write('00\n')
                     zeroes = 0
                 f.write(f'{byte:02x}\n')
+
+
+DEFAULT_WRITER = write_binary_image
+
+IMAGE_WRITERS: dict[str, WriterFn] = {
+    '.bin': write_binary_image,
+    '.img': write_logisim_image,
+}
+
+
+def write_image(filename: Union[Path, str], arr: bytearray):
+    """
+    Write the contents of array into file in format
+    that is derived from extension.
+    If extension is unknown then format is defaulted
+    to binary.
+
+    :param filename: path to output file
+    :param arr: bytearray to be written
+    """
+
+    if not isinstance(filename, Path):
+        filename = Path(filename)
+
+    writer = IMAGE_WRITERS.get(filename.suffix, DEFAULT_WRITER)
+
+    writer(filename, arr)