Implement syntax highlighting

shape/shape.py

@@ -1,5 +1,31 @@
 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 = []
@@ -15,11 +41,18 @@ def main():
     parser.add_argument("infile",
             help="input source file")
     parser.add_argument("outfile",
-            help="output image file (image format detected based on extension)")
+            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:
@@ -29,13 +62,32 @@ def main():
     width = max(map(len, lines))
     height = len(lines)
 
-    image = PIL.Image.new("RGB", (width, height), (255, 255, 255))
+    image = PIL.Image.new("RGB", (width, height),
+            PIL.ImageColor.getrgb("#" + args.bgcolor))
 
-    # Algorithm for drawing the "shape" of code, to be improved later
-    for y, line in enumerate(lines):
-        for x, char in enumerate(line):
-            if not char.isspace():
-                image.putpixel((x, y), (0, 0, 0))
+    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),