Allow saving of views as an image

.gitignore

@@ -29,5 +29,6 @@
 # Debug artifacts
 *.jsonl
 
-# Output webpages
+# Output images and webpages
 *.html
+*.png

TODO.md

@@ -23,7 +23,7 @@ Wish list items in a future version:
 - [ ] Add a Google Storage abstraction
 - [ ] Add an Azure Blob abstraction
 - [ ] Add a FTP abstraction
-- [ ] A local view of some sort
+- [x] A local view of some sort
 - [ ] Options to control the main size in the HTML view 
 - [ ] Some level of integration with CloudWatch as a widget
 - [ ] An example Lambda that auto-generates a webpage with some frequency

dir_sizer.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 from datetime import datetime
-from grid_layout import get_webpage, AUTO_SCALE, SET_SIZE
+from grid_layout import get_webpage, get_image, AUTO_SCALE, SET_SIZE
 from utils import Folder, BatchingSql, ALL_ABSTRACTIONS
 import json
 import os
@@ -29,6 +29,21 @@ def set_output(opts, args):
         opts['show_help'] = True
         return args
 
+def set_output_image(opts, args):
+    if len(args) > 0:
+        if opts['output_mode'] is not None:
+            print("ERROR: Output already specified")
+            opts['show_help'] = True
+            return args
+        opts['output'] = args[0]
+        opts['output_mode'] = 'png'
+        return args[1:]
+    else:
+        print("ERROR: No filename for --output_image specified")
+        opts['show_help'] = True
+        return args
+
+
 def set_cache(opts, args):
     if len(args) > 0:
         if opts['cache'] is not None:
@@ -96,6 +111,7 @@ def main():
 
     flags = {
         '--output': set_output,
+        '--output_image': set_output_image,
         '--no_output': set_no_output,
         '--cache': set_cache,
         '--cache_dir': set_cache_dir,
@@ -155,9 +171,10 @@ def main():
         print(textwrap.dedent("""
             Usage: 
 
-            --output <value> = Filename to output results to
-            --no_output      = Don't create any output file
-            --debug          = Show some additional options useful for debugging
+            --output <value>       = Filename to output results to
+            --output_image <value> = Filename to output image to
+            --no_output            = Don't create any output file
+            --debug                = Show some additional options useful for debugging
         """))
         for cur in ALL_ABSTRACTIONS:
             print(f"{cur.MAIN_SWITCH} = {cur.DESCRIPTION}")
@@ -176,12 +193,13 @@ def main():
 
     if opts['output_mode'] == 'html':
         with open(opts['output'], "wt", newline="\n", encoding="utf-8") as f:
-            # TODO: Let the final size be an option
             # The values 1900x965 are designed to be about the real-estate available
             # on a 1080p display with some space left over for the browser UI
             f.write(get_webpage(opts, abstraction, folder, 1900, 965, AUTO_SCALE))
-            # A test set size
-            # f.write(get_webpage(opts, abstraction, folder, 900, 600, SET_SIZE))
+        print(f"All done, created {opts['output']}")
+    elif opts['output_mode'] == 'png':
+        im = get_image(opts, abstraction, folder, 1920, 1080)
+        im.save(opts['output'])
         print(f"All done, created {opts['output']}")
     elif opts['output_mode'] == 'none':
         print(f"All done")

grid_layout.py

@@ -153,7 +153,7 @@ def _get_width(temp):
         temp.vertical = True
         return temp.height
 
-def get_color(depth):
+def get_color(depth, as_rgb=False):
     # Return a color for a grid cell, only the depth of the cell is taken into context
     depth += 1
     hue = depth * 40
@@ -170,10 +170,18 @@ def get_color(depth):
         green = green * saturation + 1.0 - saturation;
         blue = blue * saturation + 1.0 - saturation;
 
-    # Return the color as an HTML color
-    return f"#{int(red*255):02x}{int(green*255):02x}{int(blue*255):02x}"
+    red = max(0, min(255, int(red * 255)))
+    green = max(0, min(255, int(green * 255)))
+    blue = max(0, min(255, int(blue * 255)))
 
-def draw_layout(opts, abstraction, width, height, x, y, scale_mode, folder, tooltips, path, depth=0, other=None):
+    if as_rgb:
+        # Return a RGB tuple
+        return red, green, blue
+    else:
+        # Return the color as an HTML color
+        return f"#{red:02x}{green:02x}{blue:02x}"
+
+def draw_layout_html(opts, abstraction, width, height, x, y, scale_mode, folder, tooltips, path, depth=0, other=None):
     # Draw a layout, returning the new layout as HTML
     if width < 20 or height < 20:
         # This cell is too small to care about
@@ -250,7 +258,7 @@ def draw_layout(opts, abstraction, width, height, x, y, scale_mode, folder, tool
     for cur in temp:
         # For all the sub cells, show them if they're big enough
         if cur.width >= 10 and cur.height >= 10 and cur.key is not None:
-            output_html += draw_layout(
+            output_html += draw_layout_html(
                 opts,
                 abstraction, 
                 cur.width, cur.height, cur.x, cur.y, scale_mode,
@@ -269,6 +277,73 @@ def draw_layout(opts, abstraction, width, height, x, y, scale_mode, folder, tool
 
     return output_html
 
+def draw_layout_image(dr, fnt, opts, abstraction, width, height, x, y, folder, path, depth=0, other=None):
+    # Draw a layout into an image
+    if width < 20 or height < 20:
+        # This cell is too small to care about
+        return
+
+    # Make this cell's width/height visible to the children so they 
+    # can size the grid based off the idealized size
+    next_other = {
+        "x": x,
+        "y": y,
+        "width": width,
+        "height": height,
+    }
+
+    # And draw each cell
+    if depth > 0:
+        dr.rounded_rectangle(
+            (x, y, x+width, y+height), 
+            fill=get_color(depth, True), 
+            outline=(0, 0, 0), 
+            width=2, 
+            radius=5,
+        )
+
+    # Remove the offset for the border, because we'll have less space
+    width -= 2
+    height -= 2
+
+    temp = plot(width, height, folder)
+
+    # Ensure all the child elements have padding
+    padding = 5
+    base_x, base_y = x, y
+    for cur in temp:
+        x, y, right, bottom = cur.x, cur.y, cur.x + cur.width, cur.y + cur.height
+        x += padding
+        y += padding
+        right -= padding
+        bottom -= padding
+        cur.x, cur.y, cur.width, cur.height = x + base_x, y + base_y, right - x, bottom - y
+
+    for cur in temp:
+        # For all the sub cells, show them if they're big enough
+        if cur.width >= 10 and cur.height >= 10 and cur.key is not None:
+            draw_layout_image(
+                dr,
+                fnt,
+                opts,
+                abstraction, 
+                cur.width, cur.height, cur.x, cur.y,
+                folder[cur.key], 
+                path + [cur.key], 
+                depth=depth + 1,
+                other=next_other,
+            )
+
+    if depth == 1 and path[-1] is not None:
+        # The first level off the root gets a label
+        size = fnt.getsize(path[-1])
+        dr.rectangle(
+            (base_x + width - (size[0] + 2), base_y + height - (size[1] + 2), base_x + width + 2, base_y + height + 2),
+            outline=(32, 32, 32),
+            fill=(200, 200, 200),
+        )
+        dr.text((base_x + width - size[0], base_y + height - size[1]), path[-1], font=fnt, fill=(0, 0, 0))
+
 def get_summary(opts, abstraction, folder):
     # Turn the dict summary from the abstraction layer into a simple HTML header
     info = abstraction.get_summary(opts, folder)
@@ -284,7 +359,7 @@ def get_webpage(opts, abstraction, folder, width, height, scale_mode):
     # Layout everything and output to HTML
     tooltips = {}
     # Create the main HTML, along the tooltips
-    tree_html = draw_layout(opts, abstraction, width, height, 0, 0, scale_mode, folder, tooltips, [])
+    tree_html = draw_layout_html(opts, abstraction, width, height, 0, 0, scale_mode, folder, tooltips, [])
     # Create the header HTML
     summary_html, location = get_summary(opts, abstraction, folder)
     # Pull out the page title
@@ -311,5 +386,17 @@ def get_webpage(opts, abstraction, folder, width, height, scale_mode):
     # And fill in the template items
     return re.sub("{{(?P<var>[a-z_]+?)}}", lambda m: vars[m.group('var')], template)
 
+def get_image(opts, abstraction, folder, width, height):
+    # Layout everything and output to an Image
+    from PIL import Image, ImageDraw, ImageFont
+
+    im = Image.new('RGBA', (width, height), (0, 0, 0, 255))
+    dr = ImageDraw.Draw(im)
+    fnt = ImageFont.truetype(os.path.join(os.path.split(__file__)[0], "OpenSans-Regular.ttf"), 15)
+
+    draw_layout_image(dr, fnt, opts, abstraction, width, height, 0, 0, folder, [])
+
+    return im
+
 if __name__ == "__main__":
     print("This module is not meant to be run directly")