OOP-ify this sucker

Colorization can be done using a class now. This allows for more
pluggable coloring methods and lets me try out new things easier.

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>

colorhash/__main__.py

@@ -1,57 +1,14 @@
 import hashlib
 import sys
 
+from .hash_colorizer import PaletteColorizer
+
 
 # TODO - WASM compile for embedding directly in HTML
 # TODO - command line parsing for hash type, infile, forcing a palette, etc
 # TODO - file streaming for infile so we aren't loading e.g. 4GB into memory unnecessarily
-
+# TODO - option to add a caption based on the filename
-
+# TODO - palettes defined by JSON
-def hash2matrix(data: bytes, w: int, h: int) -> list[list[int]]:
-    """
-    Convert a hash to a list of rows of nibbles.
-    """
-    # Split the data into nibbles first, in case the width is an odd number
-    nibbles = []
-    for b in data:
-        top = (b & 0xF0) >> 4
-        bottom = b & 0x0F
-        nibbles += [top, bottom]
-
-    if len(nibbles) != w * h:
-        raise ValueError(
-            f"input data length ({len(nibbles)}) must match matrix dimensions ({w}x{h} = {w * h})"
-        )
-
-    cols = []
-    row = []
-    for b in nibbles:
-        row += [b]
-        if len(row) == w:
-            cols += [row]
-            row = []
-
-    return cols
-
-
-def gensvg(matrix: list[list[int]], square_size: int = 32):
-    h = len(matrix)
-    w = len(matrix[0])
-
-    # Start SVG string
-    svg = f'<svg width="{w * square_size}" height="{h * square_size}" xmlns="http://www.w3.org/2000/svg">\n'
-
-    # Generate grid
-    for r in range(h):
-        for c in range(w):
-            x = c * square_size
-            y = r * square_size
-            color = matrix[r][c]
-            svg += f'  <rect x="{x}" y="{y}" width="{square_size}" height="{square_size}" fill="{color}" />\n'
-
-    # Close SVG string
-    svg += "</svg>"
-    return svg
 
 
 color_table = [
@@ -87,7 +44,6 @@ hash_algo = "sha512"
 
 infile = sys.stdin.buffer
 
-
 if len(sys.argv) > 1:
     inpath = sys.argv[1]
     if inpath != '-':
@@ -96,19 +52,14 @@ if len(sys.argv) > 1:
 
 hashdata = hashlib.file_digest(infile, hash_algo).digest()
 w, h = dimensions_table[hash_algo]
-matrix = hash2matrix(hashdata, w, h)
-
-# Print matrix
-# pprint.pprint([[hex(c) for c in row] for row in matrix])
 
 palette_no = sum(hashdata) % 8
-# print('using palette:', palette_no)
 palette = color_table[palette_no]
 
-colors = [[palette[v] for v in row] for row in matrix]
+colorizer = PaletteColorizer(palette)
 
 # Print colors
 # pprint.pprint([[hex(c) for c in row] for row in colors])
 
 # Print SVG
-print(gensvg(colors))
+print(colorizer.hash_to_svg(hashdata, w, h, 32))

colorhash/hash_colorizer.py

@@ -0,0 +1,78 @@
+import abc
+from typing import Sequence
+
+
+Matrix = Sequence[Sequence[int]]
+
+
+class HashColorizer(metaclass=abc.ABCMeta):
+    def hash_to_matrix(self, data: bytes, w: int, h: int) -> Matrix:
+        """
+        Convert a set of bytes to a list of rows of nibbles.
+        """
+
+        nibbles = []
+        for b in data:
+            top = (b & 0xF0) >> 4
+            bottom = b & 0x0F
+            nibbles += [top, bottom]
+
+        if len(nibbles) != w * h:
+            raise ValueError(
+                f"input data length ({len(nibbles)}) must match matrix dimensions ({w}x{h} = {w * h})"
+            )
+
+        cols = []
+        row = []
+        for b in nibbles:
+            row += [b]
+            if len(row) == w:
+                cols += [row]
+                row = []
+
+        return cols
+
+    @abc.abstractmethod
+    def colorize(self, matrix: Matrix) -> Matrix:
+        """
+        Colorize a matrix.
+        """
+
+    def gensvg(self, matrix: Matrix, square_size: int) -> str:
+        """
+        Generate an SVG based on a given matrix.
+        """
+        h = len(matrix)
+        w = len(matrix[0])
+
+        # Start SVG string
+        svg = f'<svg width="{w * square_size}" height="{h * square_size}" xmlns="http://www.w3.org/2000/svg">\n'
+
+        # Generate grid
+        for r in range(h):
+            for c in range(w):
+                x = c * square_size
+                y = r * square_size
+                color = matrix[r][c]
+                svg += f'  <rect x="{x}" y="{y}" width="{square_size}" height="{square_size}" fill="{color}" />\n'
+
+        # Close SVG string
+        svg += "</svg>"
+        return svg
+
+    def hash_to_svg(self, hash: bytes, w: int, h: int, square_size: int) -> str:
+        matrix = self.hash_to_matrix(hash, w, h)
+        colors = self.colorize(matrix)
+        return self.gensvg(colors, square_size)
+
+
+Palette = list
+
+
+class PaletteColorizer(HashColorizer):
+    def __init__(self, palette: Palette) -> None:
+        assert len(palette) == 16, "palette must contain exactly 16 colors"
+        self.palette = palette
+
+    def colorize(self, matrix: Matrix) -> Matrix:
+        return [[self.palette[v] for v in row] for row in matrix]