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 <alekratz@gmail.com>
This commit is contained in:
2024-06-04 21:58:31 -07:00
parent 1758d477a8
commit 88ceb5e2ff
3 changed files with 145 additions and 52 deletions

View File

@@ -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)

View File

@@ -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 += "</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

View File

@@ -1,42 +1,42 @@
<svg width="256" height="160" xmlns="http://www.w3.org/2000/svg">
<rect x="0" y="0" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="32" y="0" width="32" height="32" fill="hsl(120.00,100.00%,70.00%)" />
<rect x="64" y="0" width="32" height="32" fill="hsl(120.00,100.00%,53.33%)" />
<rect x="96" y="0" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="128" y="0" width="32" height="32" fill="hsl(120.00,100.00%,83.33%)" />
<rect x="160" y="0" width="32" height="32" fill="hsl(120.00,100.00%,50.00%)" />
<rect x="192" y="0" width="32" height="32" fill="hsl(120.00,100.00%,56.67%)" />
<rect x="224" y="0" width="32" height="32" fill="hsl(120.00,100.00%,63.33%)" />
<rect x="0" y="32" width="32" height="32" fill="hsl(120.00,100.00%,90.00%)" />
<rect x="32" y="32" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="64" y="32" width="32" height="32" fill="hsl(120.00,100.00%,63.33%)" />
<rect x="96" y="32" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="128" y="32" width="32" height="32" fill="hsl(120.00,100.00%,86.67%)" />
<rect x="160" y="32" width="32" height="32" fill="hsl(120.00,100.00%,60.00%)" />
<rect x="192" y="32" width="32" height="32" fill="hsl(120.00,100.00%,96.67%)" />
<rect x="224" y="32" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="0" y="64" width="32" height="32" fill="hsl(120.00,100.00%,56.67%)" />
<rect x="32" y="64" width="32" height="32" fill="hsl(120.00,100.00%,63.33%)" />
<rect x="64" y="64" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="96" y="64" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="128" y="64" width="32" height="32" fill="hsl(120.00,100.00%,56.67%)" />
<rect x="160" y="64" width="32" height="32" fill="hsl(120.00,100.00%,96.67%)" />
<rect x="192" y="64" width="32" height="32" fill="hsl(120.00,100.00%,73.33%)" />
<rect x="224" y="64" width="32" height="32" fill="hsl(120.00,100.00%,53.33%)" />
<rect x="0" y="96" width="32" height="32" fill="hsl(120.00,100.00%,86.67%)" />
<rect x="32" y="96" width="32" height="32" fill="hsl(120.00,100.00%,76.67%)" />
<rect x="64" y="96" width="32" height="32" fill="hsl(120.00,100.00%,90.00%)" />
<rect x="96" y="96" width="32" height="32" fill="hsl(120.00,100.00%,60.00%)" />
<rect x="128" y="96" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="160" y="96" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="192" y="96" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="224" y="96" width="32" height="32" fill="hsl(120.00,100.00%,80.00%)" />
<rect x="0" y="128" width="32" height="32" fill="hsl(120.00,100.00%,100.00%)" />
<rect x="32" y="128" width="32" height="32" fill="hsl(120.00,100.00%,66.67%)" />
<rect x="64" y="128" width="32" height="32" fill="hsl(120.00,100.00%,70.00%)" />
<rect x="96" y="128" width="32" height="32" fill="hsl(120.00,100.00%,70.00%)" />
<rect x="128" y="128" width="32" height="32" fill="hsl(120.00,100.00%,66.67%)" />
<rect x="160" y="128" width="32" height="32" fill="hsl(120.00,100.00%,70.00%)" />
<rect x="192" y="128" width="32" height="32" fill="hsl(120.00,100.00%,83.33%)" />
<rect x="224" y="128" width="32" height="32" fill="hsl(120.00,100.00%,93.33%)" />
<rect x="0" y="0" width="32" height="32" fill="hsl(180.00,100.00%,53.33%)" />
<rect x="32" y="0" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="64" y="0" width="32" height="32" fill="hsl(180.00,100.00%,66.67%)" />
<rect x="96" y="0" width="32" height="32" fill="hsl(180.00,100.00%,76.67%)" />
<rect x="128" y="0" width="32" height="32" fill="hsl(180.00,100.00%,93.33%)" />
<rect x="160" y="0" width="32" height="32" fill="hsl(180.00,100.00%,63.33%)" />
<rect x="192" y="0" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="224" y="0" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="0" y="32" width="32" height="32" fill="hsl(180.00,100.00%,83.33%)" />
<rect x="32" y="32" width="32" height="32" fill="hsl(180.00,100.00%,76.67%)" />
<rect x="64" y="32" width="32" height="32" fill="hsl(180.00,100.00%,66.67%)" />
<rect x="96" y="32" width="32" height="32" fill="hsl(180.00,100.00%,96.67%)" />
<rect x="128" y="32" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="160" y="32" width="32" height="32" fill="hsl(180.00,100.00%,86.67%)" />
<rect x="192" y="32" width="32" height="32" fill="hsl(180.00,100.00%,83.33%)" />
<rect x="224" y="32" width="32" height="32" fill="hsl(180.00,100.00%,66.67%)" />
<rect x="0" y="64" width="32" height="32" fill="hsl(180.00,100.00%,66.67%)" />
<rect x="32" y="64" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="64" y="64" width="32" height="32" fill="hsl(180.00,100.00%,66.67%)" />
<rect x="96" y="64" width="32" height="32" fill="hsl(180.00,100.00%,63.33%)" />
<rect x="128" y="64" width="32" height="32" fill="hsl(180.00,100.00%,80.00%)" />
<rect x="160" y="64" width="32" height="32" fill="hsl(180.00,100.00%,93.33%)" />
<rect x="192" y="64" width="32" height="32" fill="hsl(180.00,100.00%,60.00%)" />
<rect x="224" y="64" width="32" height="32" fill="hsl(180.00,100.00%,60.00%)" />
<rect x="0" y="96" width="32" height="32" fill="hsl(180.00,100.00%,90.00%)" />
<rect x="32" y="96" width="32" height="32" fill="hsl(180.00,100.00%,96.67%)" />
<rect x="64" y="96" width="32" height="32" fill="hsl(180.00,100.00%,83.33%)" />
<rect x="96" y="96" width="32" height="32" fill="hsl(180.00,100.00%,80.00%)" />
<rect x="128" y="96" width="32" height="32" fill="hsl(180.00,100.00%,83.33%)" />
<rect x="160" y="96" width="32" height="32" fill="hsl(180.00,100.00%,63.33%)" />
<rect x="192" y="96" width="32" height="32" fill="hsl(180.00,100.00%,73.33%)" />
<rect x="224" y="96" width="32" height="32" fill="hsl(180.00,100.00%,50.00%)" />
<rect x="0" y="128" width="32" height="32" fill="hsl(180.00,100.00%,70.00%)" />
<rect x="32" y="128" width="32" height="32" fill="hsl(180.00,100.00%,70.00%)" />
<rect x="64" y="128" width="32" height="32" fill="hsl(180.00,100.00%,90.00%)" />
<rect x="96" y="128" width="32" height="32" fill="hsl(180.00,100.00%,50.00%)" />
<rect x="128" y="128" width="32" height="32" fill="hsl(180.00,100.00%,53.33%)" />
<rect x="160" y="128" width="32" height="32" fill="hsl(180.00,100.00%,86.67%)" />
<rect x="192" y="128" width="32" height="32" fill="hsl(180.00,100.00%,60.00%)" />
<rect x="224" y="128" width="32" height="32" fill="hsl(180.00,100.00%,83.33%)" />
</svg>

Before

Width:  |  Height:  |  Size: 3.3 KiB

After

Width:  |  Height:  |  Size: 3.3 KiB