Fix types for mypy

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

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",

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]

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.
     """