file-visualisation/shape/shape.py

100 lines
3.1 KiB
Python

import argparse
import json
import PIL.Image
import PIL.ImageColor
import pygments
import pygments.formatter
import pygments.lexers
class JsonFormatter(pygments.formatter.Formatter):
def __init__(self, **options):
super().__init__(**options)
# Probably unnecessary, but hey, it works (hopefully)
self.styles = {}
for token, style in self.style:
info = {"color": style["color"], "bgcolor": style["bgcolor"]}
self.styles[token] = info
def format(self, tokensource, outfile):
tokens = []
for ttype, value in tokensource:
while ttype not in self.styles:
ttype = ttype.parent
info = {"value": value, "style": self.styles[ttype]}
tokens.append(info)
outfile.write(json.dumps(tokens))
def convert_tabs(line, tabwidth):
result = []
for char in line:
if char == "\t":
result.append(" " * tabwidth)
else:
result.append(char)
return "".join(result)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("infile",
help="input source file")
parser.add_argument("outfile",
help=("output image file (image format detected based on"
" extension)"))
parser.add_argument("--tabwidth", "-w", type=int, default=8,
help="the amount of spaces per tab")
parser.add_argument("--upscale", "-u", action="store_true",
help="increase the output image's size by 10")
parser.add_argument("--textcolor", "-t", default="000000",
help=("default text color for all sections that pygment doesn't"
" specify a color for"))
parser.add_argument("--bgcolor", "-b", default="FFFFFF",
help=("default background color for all sections that pygment"
" doesn't specify a color for"))
args = parser.parse_args()
with open(args.infile) as f:
text = f.read()
lines = [convert_tabs(line, args.tabwidth) for line in text.splitlines()]
width = max(map(len, lines))
height = len(lines)
image = PIL.Image.new("RGB", (width, height),
PIL.ImageColor.getrgb("#" + args.bgcolor))
tokens = json.loads(pygments.highlight(text, pygments.lexers.PythonLexer(),
JsonFormatter()))
x, y = 0, 0
for token in tokens:
value, style = token["value"], token["style"]
colorstr = "#" + (style["color"] or args.textcolor)
color = PIL.ImageColor.getrgb(colorstr)
bgcolorstr = "#" + (style["bgcolor"] or args.bgcolor)
bgcolor = PIL.ImageColor.getrgb(bgcolorstr)
for char in value:
if char == "\n":
x = 0
y += 1
elif char.isspace():
image.putpixel((x, y), bgcolor)
x += 1
else:
image.putpixel((x, y), color)
x += 1
if args.upscale:
image = image.resize((width * 10, height * 10),
resample=PIL.Image.NEAREST)
with open(args.outfile, "wb") as f:
image.save(f)
if __name__ == "__main__":
main()