From 1277c2db8bf6805116d60cf621e124224a123c63 Mon Sep 17 00:00:00 2001 From: Alek Ratzloff Date: Fri, 24 May 2024 11:45:42 -0700 Subject: [PATCH] Add input-type Input can now either be a path, raw data, or a hash itself (path is default) Signed-off-by: Alek Ratzloff --- colorhash/__main__.py | 71 +++++++++++++++++++++++++++++++++--------- colorhash/colorizer.py | 5 +-- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/colorhash/__main__.py b/colorhash/__main__.py index 06b05c2..5641338 100644 --- a/colorhash/__main__.py +++ b/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 - # file_digest (I hope) will not load too much into memory - hashdata = hashlib.file_digest(args.infile, args.hash).digest() + 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(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) diff --git a/colorhash/colorizer.py b/colorhash/colorizer.py index e41352f..970c34b 100644 --- a/colorhash/colorizer.py +++ b/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: