diff --git a/shape/shape.py b/shape/shape.py index 1899c25..5139287 100644 --- a/shape/shape.py +++ b/shape/shape.py @@ -1,5 +1,6 @@ import argparse import json +import sys import PIL.Image import PIL.ImageColor @@ -27,16 +28,107 @@ class JsonFormatter(pygments.formatter.Formatter): tokens.append(info) outfile.write(json.dumps(tokens)) -def convert_tabs(line, tabwidth): - result = [] - for char in line: - if char == "\t": - result.append(" " * tabwidth) +class Shaper: + def __init__(self, args): + self.args = args + + def convert_tabs(self, string): + return string.replace("\t", " " * self.args.tabwidth) + + def get_dimensions(self, text): + lines = self.convert_tabs(text).splitlines() + width = max(map(len, lines)) + height = len(lines) + return width, height + + def to_color(self, colorstr): + return PIL.ImageColor.getrgb("#" + colorstr) + + def figure_out_lexer(self, text): + if self.args.lexer is not None: + return pygments.lexers.get_lexer_by_name(self.args.lexer) else: - result.append(char) - return "".join(result) + return pygments.lexers.guess_lexer_for_filename(self.args.infile, text) + + def get_style(self): + if self.args.style is not None: + return pygments.styles.get_style_by_name(self.args.style) + + def draw_to_image(self, image, tokens): + x, y = 0, 0 + for token in tokens: + value, style = token["value"], token["style"] + color = self.to_color(style["color"] or self.args.textcolor) + bgcolor = self.to_color(style["bgcolor"] or self.args.bgcolor) + + 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 + + def find_shape(self): + with open(self.args.infile) as f: + text = f.read() + + width, height = self.get_dimensions(text) + print(f"Image dimensions: {width}x{height}") + image = PIL.Image.new( "RGB", (width, height), + self.to_color(self.args.bgcolor)) + + lexer = self.figure_out_lexer(text) + print(f"Using lexer: {lexer.name}") + style = self.get_style() + if style is None: + print("Using default style") + formatter = JsonFormatter() + else: + print(f"Using style: {self.args.style}") + formatter = JsonFormatter(style=style) + tokens = json.loads(pygments.highlight(text, lexer, formatter)) + + self.draw_to_image(image, tokens) + + if self.args.upscale: + print("Scaling up the image by a factor of 10") + image = image.resize((width * 10, height * 10), + resample=PIL.Image.NEAREST) + print(f"New image dimensions: {image.width}x{image.height}") + + with open(self.args.outfile, "wb") as f: + image.save(f) + +def list_lexers(): + print("Available lexers:") + + short_names = [info[1] for info in pygments.lexers.get_all_lexers()] + for names in sorted(short_names): + print(" " + ", ".join(names)) + +def list_styles(): + print("Available styles:") + + for style in sorted(pygments.styles.get_all_styles()): + print(" " + style) def main(): + # Workaround because I don't know how to do this with argparse + lexers = "--list-lexers" in sys.argv + styles = "--list-styles" in sys.argv + if lexers or styles: + if lexers: + list_lexers() + if lexers and styles: + print() + if styles: + list_styles() + return + parser = argparse.ArgumentParser() parser.add_argument("infile", help="input source file") @@ -53,48 +145,18 @@ def main(): parser.add_argument("--bgcolor", "-b", default="FFFFFF", help=("default background color for all sections that pygment" " doesn't specify a color for")) + parser.add_argument("--lexer", "-l", + help="the lexer to use (for a list, see --list-lexers)") + parser.add_argument("--list-lexers", action="store_true", + help="a list of all lexers available") + parser.add_argument("--style", "-s", + help="the color scheme to use (for a list, see --list-styles)") + parser.add_argument("--list-styles", action="store_true", + help="a list of all color schemes available") + 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) + shaper = Shaper(args) + shaper.find_shape() if __name__ == "__main__": main()