Support conversion directly from NumPy arrays (#33)

Support conversion directly from NumPy arrays

Author: gahjelle

.github/workflows/test.yml

@@ -24,7 +24,7 @@ jobs:
                   pip install flit pytest
             - name: Install module
               run: |
-                  flit install
+                  flit install --extras test
             - name: Run tests
               run: |
                   pytest
@@ -44,7 +44,7 @@ jobs:
                   pip install flit pytest coverage
             - name: Install module
               run: |
-                  flit install
+                  flit install --extras test
             - name: Run tests
               run: |
                   coverage run -m pytest

README.md

@@ -1,6 +1,6 @@
 # Advent of Code® OCR
 
-This Python module helps with converting Advent of Code ASCII art letters into plain characters. At the moment, it only supports 6-pixel-tall characters as seen in 2016 Day 8, 2019 Days 8 and 11, and 2021 Day 13.
+This Python module helps with converting [Advent of Code](https://adventofcode.com/) ASCII art letters into plain characters. At the moment, it only supports 6-pixel-tall characters as seen in 2016 Day 8, 2019 Days 8 and 11, and 2021 Day 13.
 
 Support for 10-pixel-tall characters (2018 Day 10) is coming soon.
 
@@ -37,6 +37,23 @@ print(convert_6(" $$ \n$  $\n$  $\n$$$$\n$  $\n$  $", fill_pixel="$", empty_pixe
 # A
 ```
 
+You can also convert data that you have in a NumPy array or a nested list:
+
+```py
+from advent_of_code_ocr import convert_array_6
+
+array = [
+    [0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0],
+    [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1],
+    [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0],
+    [1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0],
+    [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1],
+    [1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0],
+]
+print(convert_array_6(array, fill_pixel=1, empty_pixel=0))
+# AOC
+```
+
 ---
 
 Advent of Code is a registered trademark of Eric K Wastl in the United States.

advent_of_code_ocr/__init__.py

@@ -1,5 +1,9 @@
 """Convert Advent of Code ASCII art"""
 
+from __future__ import annotations
+
+from collections.abc import Sequence
+
 from .characters import ALPHABET_6
 
 __version__ = "1.0.0"
@@ -8,9 +12,31 @@
 def convert_6(
     input_text: str, *, fill_pixel: str = "#", empty_pixel: str = "."
 ) -> str:
-    """Convert height 6 characters"""
+    """Convert height 6 text to characters"""
     input_text = input_text.replace(fill_pixel, "#").replace(empty_pixel, ".")
-    array = input_text.split("\n")
+    prepared_array = [list(line) for line in input_text.split("\n")]
+    return _convert_6(prepared_array)
+
+
+def convert_array_6(
+    array: Sequence[Sequence[str | int]],
+    *,
+    fill_pixel: str | int = "#",
+    empty_pixel: str | int = "."
+) -> str:
+    """Convert a height 6 NumPy array or nested list to characters"""
+    prepared_array = [
+        [
+            "#" if pixel == fill_pixel else "." if pixel == empty_pixel else ""
+            for pixel in line
+        ]
+        for line in array
+    ]
+    return _convert_6(prepared_array)
+
+
+def _convert_6(array: list[list[str]]) -> str:
+    """Convert a prepared height 6 array to characters"""
 
     # Validate input
     rows, cols = len(array), len(array[0])
@@ -22,7 +48,8 @@ def convert_6(
     # Convert each letter
     indices = [slice(start, start + 4) for start in range(0, cols, 5)]
     result = [
-        ALPHABET_6["\n".join(row[index] for row in array)] for index in indices
+        ALPHABET_6["\n".join("".join(row[index]) for row in array)]
+        for index in indices
     ]
 
     return "".join(result)

pyproject.toml

@@ -18,7 +18,6 @@ classifiers = [
   "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
 ]
-dependencies = []
 dynamic = ["version", "description"]
 license = {file = "LICENSE"}
 name = "advent-of-code-ocr"
@@ -32,6 +31,9 @@ name = "advent_of_code_ocr"
 Changelog = "https://github.com/bsoyka/advent-of-code-ocr/releases"
 Source = "https://github.com/bsoyka/advent-of-code-ocr"
 
+[project.optional-dependencies]
+test = ["numpy", "pytest"]
+
 ### TOOLS ###
 
 [tool.black]
@@ -65,6 +67,7 @@ envlist = py37,py38,py39,py310
 isolated_build = True
 
 [testenv]
-deps = pytest
+deps = .
 commands = pytest
+extras = test
 """

tests/test_6.py

@@ -1,8 +1,9 @@
 from re import escape
 
+import numpy as np
 from pytest import mark, raises
 
-from advent_of_code_ocr import convert_6
+from advent_of_code_ocr import convert_6, convert_array_6
 from advent_of_code_ocr.characters import ALPHABET_6
 
 
@@ -82,3 +83,49 @@ def test_strange_width_characters():
         ]
     )
     assert convert_6(string) == "EFEYKFRFIJ"
+
+
+def test_array_nested_list():
+    array = [
+        ["X", "O", "O", "X", "X", "X", "O", "O", "X"],
+        ["O", "X", "X", "O", "X", "O", "X", "X", "O"],
+        ["O", "X", "X", "O", "X", "O", "X", "X", "X"],
+        ["O", "O", "O", "O", "X", "O", "X", "X", "X"],
+        ["O", "X", "X", "O", "X", "O", "X", "X", "O"],
+        ["O", "X", "X", "O", "X", "X", "O", "O", "X"],
+    ]
+    assert convert_array_6(array, fill_pixel="O", empty_pixel="X") == "AC"
+
+
+def test_array_list_of_strings():
+    array = [
+        ".oo...oo..oooo..ooo.o..o..ooo..oo..o..o",
+        "o..o.o..o.o....o....o..o...o..o..o.o..o",
+        "o..o.o....ooo..o....oooo...o..o....oooo",
+        "oooo.o....o.....oo..o..o...o..o.oo.o..o",
+        "o..o.o..o.o.......o.o..o...o..o..o.o..o",
+        "o..o..oo..oooo.ooo..o..o..ooo..ooo.o..o",
+    ]
+    assert convert_array_6(array, fill_pixel="o") == "ACESHIGH"
+
+
+def test_array_numpy():
+    array = np.array(
+        [
+            [0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0],
+            [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1],
+            [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0],
+            [1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0],
+            [1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1],
+            [1, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0],
+        ]
+    )
+    assert convert_array_6(array, fill_pixel=1, empty_pixel=0) == "AOC"
+
+
+@mark.parametrize("rows", [1, 5, 7, 10])
+def test_array_number_of_rows(rows):
+    with raises(
+        ValueError, match=escape("incorrect number of rows (expected 6)")
+    ):
+        convert_array_6(["" for _ in range(rows)])