Add input-type

Input can now either be a path, raw data, or a hash itself (path is
default)

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

colorhash/__main__.py

@@ -4,6 +4,7 @@ import argparse
 import hashlib
 from pathlib import Path
 import sys
+import textwrap
 
 from .colorizer import PaletteColorizer
 from .matricizer import Matricizer, NibbleMatricizer, RandomartMatricizer
@@ -46,11 +47,19 @@ def main() -> None:
         ]
     )
     HASH_CHOICES = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512"]
+    INPUT_TYPE_CHOICES = {
+        "path": "the input should be treated as a path and data is read from the path",
+        "hash": "the input should be treated as a hexadecimal hash (requires -a or --hash to be supplied)",
+        "data": "the input should be treated as raw data",
+    }
+    INPUT_TYPE_HELP = "INPUT TYPE (-x, --input-type)\n" + "\n".join(
+        [f"    {choice} - {desc}" for choice, desc in INPUT_TYPE_CHOICES.items()]
+    )
     EPILOGUE = "\n\n".join([MATRIX_HELP, PALETTE_HELP])
 
     progname: str = sys.argv[0]
-    if progname.endswith('__main__.py'):
-        progname = 'colorhash'
+    if progname.endswith("__main__.py"):
+        progname = "colorhash"
 
     ap = argparse.ArgumentParser(
         prog=progname,
@@ -60,10 +69,10 @@ def main() -> None:
     )
 
     ap.add_argument(
-        "infile",
-        type=argparse.FileType("rb"),
-        default=sys.stdin,
-        help="The input file to use. Set to '-' or blank for STDIN. default: STDIN",
+        "input",
+        type=str,
+        default="-",
+        help="The input to use. When acting as a path, set to '-' or blank for STDIN. Use -x or --input-type to control how input is treated. default: -",
     )
     ap.add_argument(
         "-o",
@@ -94,7 +103,8 @@ def main() -> None:
         "--hash",
         metavar="ALGORITHM",
         choices=HASH_CHOICES,
-        default="sha512",
+        # default="sha512",
+        required=False,
         help="Choose the hash algorithm. default: sha512",
     )
     ap.add_argument(
@@ -105,30 +115,63 @@ def main() -> None:
         default=32,
         help="Decide how big the output squares are, in pixels. default: 32",
     )
+    ap.add_argument(
+        "-x",
+        "--input-type",
+        default="path",
+        choices=INPUT_TYPE_CHOICES.keys(),
+        help="Determines how the input should be treated. default: path",
+    )
     args = ap.parse_args()
 
     ############################################################################
     # End arg parsing
     ############################################################################
 
+    # -a/--hash arg is not required when we're using file and data input types. only required for
+    # hash input type
+    if args.input_type in ("data", "path") and args.hash is None:
+        args.hash = "sha512"
+
     # Get the hash
+    match args.input_type:
+        case "path":
+            if args.input == "-":
+                infile = sys.stdin.buffer
+            else:
+                # TODO - pretty error message for when the file doesn't exist
+                infile = open(args.input, "rb")
             # file_digest (I hope) will not load too much into memory
-    hashdata = hashlib.file_digest(args.infile, args.hash).digest()
+            hashdata = hashlib.file_digest(infile, args.hash).digest()
+        case "hash":
+            # TODO - maybe a better error message?
+            if args.hash is None:
+                print(
+                    "ERROR: -a or --hash should be supplied on the command line when using the hash input type",
+                    file=sys.stderr,
+                )
+                raise SystemExit(1)
+            # TODO - pretty error message for malformed input
+            hashdata = bytes([int(byte, 16) for byte in textwrap.wrap(args.input, 2)])
+        case "data":
+            hashdata = hashlib.new(args.hash, args.input.encode()).digest()
+        case _:
+            assert False, f"unknown input type {args.input_type}"
 
     # Choose the palette
     palette: list[str]
-    if args.palette == 'auto':
-        palette = list(PALETTES.values())[sum(hashdata) % 8]
+    if args.palette == "auto":
+        palette = list(DEFAULT_PALETTES.values())[sum(hashdata) % 8]
     else:
         palette = PALETTES[args.palette]
 
     # Choose the dimensions and the matricizer
     matricizer: Matricizer
     match args.matrix:
-        case 'nibble':
+        case "nibble":
             w, h = NibbleMatricizer.DIMENSIONS[args.hash]
             matricizer = NibbleMatricizer(w, h)
-        case 'randomart':
+        case "randomart":
             # 17x9 is what openssh uses
             # TODO - allow configuring dimensions, maybe
             matricizer = RandomartMatricizer(17, 9)
@@ -142,7 +185,7 @@ def main() -> None:
     matrix = matricizer.matricize(hashdata)
     colors = colorizer.colorize(matrix)
     svg = gensvg(colors, args.square_size)
-    if str(args.out) == '-':
+    if str(args.out) == "-":
         sys.stdout.write(svg)
     else:
         args.out.write_text(svg)

colorhash/colorizer.py

@@ -3,6 +3,7 @@ import abc
 from typing import Sequence
 
 from .matricizer import Matrix
+from .palettes import Palette
 
 
 StrMatrix = Sequence[Sequence[str]]
@@ -26,9 +27,6 @@ class Colorizer(metaclass=abc.ABCMeta):
         """
 
 
-Palette = Sequence[str]
-
-
 class PaletteColorizer(Colorizer):
     """
     A palette colorizer.
@@ -41,7 +39,6 @@ class PaletteColorizer(Colorizer):
 
         :param palette: the palette to use for this colorizer.
         """
-        assert len(palette) == 16, "palette must contain exactly 16 colors"
         self.palette = palette
 
     def colorize(self, matrix: Matrix) -> StrMatrix: