From 840e878748466fb08cfefce3de864b55f55c1ee5 Mon Sep 17 00:00:00 2001
From: Alek Ratzloff <alekratz@gmail.com>
Date: Wed, 22 May 2024 11:15:43 -0700
Subject: [PATCH] Fix types for mypy

Signed-off-by: Alek Ratzloff <alekratz@gmail.com>
---
 colorhash/__main__.py  | 4 ++--
 colorhash/colorizer.py | 7 +++++--
 colorhash/svg.py       | 4 ++--
 3 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/colorhash/__main__.py b/colorhash/__main__.py
index f278d20..da78f4e 100644
--- a/colorhash/__main__.py
+++ b/colorhash/__main__.py
@@ -4,7 +4,7 @@ from pathlib import Path
 import sys
 
 from .colorizer import PaletteColorizer
-from .matricizer import NibbleMatricizer, RandomartMatricizer
+from .matricizer import Matricizer, NibbleMatricizer, RandomartMatricizer
 from .svg import gensvg
 
 
@@ -25,7 +25,7 @@ PALETTES = {
 }
 
 
-def main():
+def main() -> None:
     MATRIX_CHOICES = {
         "nibble": "Use each nibble (4 bits) of the hash to generate a matrix",
         "randomart": "Use the SSH 'randomart' algorithm to generate a matrix",
diff --git a/colorhash/colorizer.py b/colorhash/colorizer.py
index 343d364..19d6a4a 100644
--- a/colorhash/colorizer.py
+++ b/colorhash/colorizer.py
@@ -4,9 +4,12 @@ from typing import Sequence
 from .matricizer import Matrix
 
 
+StrMatrix = Sequence[Sequence[str]]
+
+
 class Colorizer(metaclass=abc.ABCMeta):
     @abc.abstractmethod
-    def colorize(self, matrix: Matrix) -> Matrix:
+    def colorize(self, matrix: Matrix) -> StrMatrix:
         """
         Colorize a matrix.
         """
@@ -20,5 +23,5 @@ class PaletteColorizer(Colorizer):
         assert len(palette) == 16, "palette must contain exactly 16 colors"
         self.palette = palette
 
-    def colorize(self, matrix: Matrix) -> Matrix:
+    def colorize(self, matrix: Matrix) -> StrMatrix:
         return [[self.palette[v] for v in row] for row in matrix]
diff --git a/colorhash/svg.py b/colorhash/svg.py
index 5b2f14a..e1c5c53 100644
--- a/colorhash/svg.py
+++ b/colorhash/svg.py
@@ -1,7 +1,7 @@
-from .matricizer import Matrix
+from .colorizer import StrMatrix
 
 
-def gensvg(matrix: Matrix, square_size: int) -> str:
+def gensvg(matrix: StrMatrix, square_size: int) -> str:
     """
     Generate an SVG based on a given matrix.
     """