From 88ceb5e2ff1490a9bb77b0edc5eae8d90745c2ff Mon Sep 17 00:00:00 2001 From: Alek Ratzloff Date: Tue, 4 Jun 2024 21:58:31 -0700 Subject: [PATCH] Add PNGWriter * png files are an available output in the CLI * --svg-square-size is now --square-size * output type of Writer.write is bytes instead of a string. Just use str.decode if string output is needed Signed-off-by: Alek Ratzloff --- colorhash/cli.py | 13 +++-- colorhash/writer.py | 104 +++++++++++++++++++++++++++++++++++++--- examples/commithash.svg | 80 +++++++++++++++---------------- 3 files changed, 145 insertions(+), 52 deletions(-) diff --git a/colorhash/cli.py b/colorhash/cli.py index a737211..73e235f 100644 --- a/colorhash/cli.py +++ b/colorhash/cli.py @@ -53,7 +53,8 @@ def cli_main() -> None: ) OUTPUT_TYPE_CHOICES = { "ansi": "the output should be colored for ANSI terminals using 24 bit true color", - "svg": "the output should be an SVG format", + "svg": "the output should be in SVG format", + "png": "the output should be in PNG format", } OUTPUT_TYPE_HELP = "OUTPUT TYPE (-y, --output-type)\n" + "\n".join( [f" {choice} - {desc}" for choice, desc in OUTPUT_TYPE_CHOICES.items()] @@ -113,7 +114,7 @@ def cli_main() -> None: help="Choose the hash algorithm. default: sha512", ) ap.add_argument( - "--svg-square-size", + "--square-size", metavar="PX", type=int, default=32, @@ -199,11 +200,13 @@ def cli_main() -> None: case "ansi": writer = ANSIWriter() case "svg": - writer = SVGWriter(args.svg_square_size) + writer = SVGWriter(args.square_size) + case "png": + writer = PNGWriter(args,square_size) output = writer.write(colors) if str(args.out) == "-": - sys.stdout.write(output) + sys.stdout.buffer.write(output) else: - args.out.write_text(output) + args.out.write_bytes(output) diff --git a/colorhash/writer.py b/colorhash/writer.py index a4e22d1..4f5c515 100644 --- a/colorhash/writer.py +++ b/colorhash/writer.py @@ -1,5 +1,6 @@ "Colorhash writer classes" import abc +import zlib from .color import Color, ColorMatrix @@ -13,12 +14,12 @@ class Writer(metaclass=abc.ABCMeta): """ @abc.abstractmethod - def write(self, matrix: ColorMatrix) -> str: + def write(self, matrix: ColorMatrix) -> bytes: """ Write the color matrix to a string. - :param matrix: the color matrix to generate the SVG for. - :returns: the full generated SVG as a string. + :param matrix: the color matrix to generate the image for. + :returns: the generated image as a string. """ @@ -27,7 +28,7 @@ class ANSIWriter(Writer): ANSI terminal writer. This will output a 24-bit true color string. """ - def write(self, matrix: ColorMatrix) -> str: + def write(self, matrix: ColorMatrix) -> bytes: """ Write the color matrix to an ANSI string. @@ -49,7 +50,7 @@ class ANSIWriter(Writer): out += c out += "\n" out += reset - return out + return out.encode() class SVGWriter(Writer): @@ -65,7 +66,7 @@ class SVGWriter(Writer): """ self.square_size = square_size - def write(self, matrix: ColorMatrix) -> str: + def write(self, matrix: ColorMatrix) -> bytes: """ Generate an SVG based on a given matrix. @@ -88,4 +89,93 @@ class SVGWriter(Writer): # Close SVG string svg += "" - return svg + return svg.encode() + + +class PNGWriter(Writer): + def __init__(self, square_size: int) -> None: + """ + Create a new PNG writer that uses the given square size. + + :param square_size: the size of the squares generated, in pixels. + """ + self.square_size = square_size + + def write(self, matrix: ColorMatrix) -> bytes: + """ + Generate a PNG based on a given matrix. + + :param matrix: the color matrix to generate the SVG for. + :returns: the full generated PNG as an ASCII-encoded string. It's probably a good idea to + convert this to bytes since it's binary data being shoved into a string type. + """ + w = self.square_size * len(matrix[0]) + h = self.square_size * len(matrix) + + def i32(i: int) -> bytes: + return int.to_bytes(i, 4, "big") + + def chunk(name: str, data: bytes) -> bytes: + assert len(name) == 4, "chunk name must be exactly 4 bytes" + chunk = bytearray() + chunk += i32(len(data)) + # add the name to the data so it also gets encoded with the crc32 + data = name.encode('ascii') + data + chunk += data + chunk += i32(zlib.crc32(data)) + return bytes(chunk) + + # Convert the matrix into RGB byte triples + colors = [ + [ + bytes([int(c.r), int(c.g), int(c.b)]) + for c in map(lambda c: c.to_rgb(), row) + ] + for row in matrix + ] + + # Create the palette based on the unique colors available + # NOTE : these could be done in the same dict and would probably save a little bit of + # memory, however, mypy likes it when we keep types easy + pal2col = {} + col2pal = {} + for i, c in enumerate(set(sum(colors, []))): + pal2col[i] = c + col2pal[c] = i + + assert len(pal2col) // 2 < 16, "palette for PNG image was longer than 16 colors" + + # Header + png = bytearray([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) + + # write the IHDR chunk + png += chunk( + "IHDR", + # width, height, bit depth (4), color type (3, palette), + # compression method (always 0), filter method (always 0), + # interlace method (0, not interlaced) + i32(w) + i32(h) + bytes([4, 3, 0, 0, 0]), + ) + # write the palette chunk + png += chunk( + "PLTE", + b"".join([pal2col[i] for i in range(len(pal2col))]), + ) + # create scanlines and shove them into IDAT chunks + idat = bytearray() + for row in colors: + line = bytearray([0]) + for col in row: + b = (col2pal[col] << 4) | col2pal[col] + # add square_size number of colors (divided by 2, since we only need 4 bits per + # color) + for _ in range(self.square_size // 2): + line += bytes([b]) + # add square_size number of lines + for _ in range(self.square_size): + idat += line + # write the IDAT chunk + png += chunk("IDAT", zlib.compress(idat)) + # write the IEND chunk + png += chunk("IEND", bytes()) + return png diff --git a/examples/commithash.svg b/examples/commithash.svg index b885610..debd3f4 100644 --- a/examples/commithash.svg +++ b/examples/commithash.svg @@ -1,42 +1,42 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file