From 65f905545218942c004dab37090ced24f39ac50f Mon Sep 17 00:00:00 2001
From: David Bau <david.bau@gmail.com>
Date: Mon, 21 Mar 2022 00:02:08 -0400
Subject: [PATCH] Rewrite show, and modify labwidget to use it.

---
 baulab/__init__.py    |  14 ++
 baulab/labwidget.py   | 125 +++++------
 baulab/paintwidget.py |  24 +--
 baulab/plotwidget.py  |   6 +-
 baulab/show.py        | 491 +++++++++++++++++++++++++++++++-----------
 5 files changed, 454 insertions(+), 206 deletions(-)
 create mode 100644 baulab/__init__.py

diff --git a/baulab/__init__.py b/baulab/__init__.py
new file mode 100644
index 0000000..06f9eee
--- /dev/null
+++ b/baulab/__init__.py
@@ -0,0 +1,14 @@
+from .labwidget import Model, Widget, Trigger, Property, Event
+from .labwidget import Button, Label, Textbox, Numberbox, Range, ColorPicker
+from .labwidget import Choice, Menu, Datalist, Div, ClickDiv, Image
+from .paintwidget import PaintWidget
+from .plotwidget import PlotWidget
+from . import pbar
+from .nethook import Trace, TraceDict, set_requires_grad
+from .nethook import subsequence, get_module, get_parameter, replace_module
+from .runningstats import Stat, Mean, Variance, Covariance, Bincount, CrossCovariance
+from .runningstats import IoU, CrossIoU, Quantile, TopK, History, CombinedStat
+from .runningstats import tally
+from . import show
+from .workerpool import WorkerBase, WorkerPool
+from .pidfile import reserve_dir
diff --git a/baulab/labwidget.py b/baulab/labwidget.py
index 5af0732..03620b9 100644
--- a/baulab/labwidget.py
+++ b/baulab/labwidget.py
@@ -49,6 +49,7 @@ import json
 import html
 import re
 from inspect import signature
+from . import show
 
 
 class Model(object):
@@ -180,7 +181,7 @@ class Widget(Model):
     for the top-level widget element.
     '''
 
-    def __init__(self, style=None, data=None, className=None):
+    def __init__(self, style=None):
         # In the jupyter case, there can be some delay between js injection
         # and comm creation, so we need to queue some initial messages.
         if WIDGET_ENV == 'jupyter':
@@ -197,8 +198,6 @@ class Widget(Model):
         # The style and data properties come standard, and are used to
         # control the style and data attributes on the toplevel element.
         self.style = Property(style)
-        self.className = Property(className)
-        self.data = Property(data)
         # Each widget has a "write" event that is used to insert
         # html before the widget.
         self.write = Trigger()
@@ -217,7 +216,8 @@ class Widget(Model):
         Override to define the initial HTML view of the widget.  Should
         define an element with id given by view_id().
         '''
-        return f'<div {self.std_attrs()}></div>'
+        with show.enter_tag() as t:
+            return t.begin() + t.end()
 
     def view_id(self):
         '''
@@ -228,13 +228,9 @@ class Widget(Model):
 
     def std_attrs(self):
         '''
-        Returns id and (if applicable) style attributes, escaped and
-        formatted for use within the top-level element of widget HTML.
+        Returns id attribute, and possibly other standard attrs.
         '''
-        return (f'id="{self.view_id()}"' +
-                style_attr(self.style) +
-                class_attr(self.className) +
-                data_attrs(self.data))
+        return show.attr(id=self.view_id())
 
     def _repr_html_(self):
         '''
@@ -557,9 +553,8 @@ class Button(Widget):
         ''')
 
     def widget_html(self):
-        return f'''<input {self.std_attrs()} type="button" value="{
-            html.escape(str(self.label))}">'''
-
+        return show.emit_tag('input', self.std_attrs(),
+                    type='button', value=self.label)
 
 class Label(Widget):
     def __init__(self, value='', **kwargs):
@@ -578,17 +573,18 @@ class Label(Widget):
         ''')
 
     def widget_html(self):
-        return f'''<label {self.std_attrs()}>{
-            html.escape(str(self.value))}</label>'''
+        out = []
+        with show.enter_tag('label', self.std_attrs(), out=out):
+            out.append(html.escape(str(self.value)))
+        return ''.join(out)
 
 
 class Textbox(Widget):
-    def __init__(self, value='', size=20, style=None, desc=None, **kwargs):
+    def __init__(self, value='', size=20, style=None, **kwargs):
         super().__init__(style=defaulted(style, display='inline-block'), **kwargs)
         # databinding is defined using Property objects.
         self.value = Property(value)
         self.size = Property(size)
-        self.desc = Property(desc)
 
     def widget_js(self):
         # Both "model" and "element" objects are defined within the scope
@@ -614,21 +610,16 @@ class Textbox(Widget):
         ''')
 
     def widget_html(self):
-
-        html_str = f'''<input {self.std_attrs()} value="{
-            html.escape(str(self.value))}" size="{self.size}">'''
-        if self.desc is not None:
-            html_str = f"""<span>{self.desc}</span>{html_str}"""
-        return html_str
+        return show.emit_tag('input', self.std_attrs(),
+                    type='text', value=self.value, size=self.size)
 
 
 class Numberbox(Widget):
-    def __init__(self, value='', size=20, style=None, desc=None, **kwargs):
+    def __init__(self, value='', size=20, style=None, **kwargs):
         super().__init__(style=defaulted(style, display='inline-block'), **kwargs)
         # databinding is defined using Property objects.
         self.value = Property(value)
         self.size = Property(size)
-        self.desc = Property(desc)
 
     def widget_js(self):
         # Both "model" and "element" objects are defined within the scope
@@ -654,12 +645,8 @@ class Numberbox(Widget):
         ''')
 
     def widget_html(self):
-
-        html_str = f'''<input {self.std_attrs()} type="numeric" value="{
-            html.escape(str(self.value))}" size="{self.size}">'''
-        if self.desc is not None:
-            html_str = f"""<span>{self.desc}</span>{html_str}"""
-        return html_str
+        return show.emit_tag('input', self.std_attrs(),
+                    type='numeric', value=self.value, size=self.size)
 
 class Range(Widget):
     def __init__(self, value=50, min=0, max=100, step=1, **kwargs):
@@ -688,8 +675,9 @@ class Range(Widget):
         ''')
 
     def widget_html(self):
-        return f'''<input {self.std_attrs()} type="range" value="{
-            self.value}" min="{self.min}" max="{self.max}" step="{self.step}">'''
+        return show.emit_tag('input', self.std_attrs(),
+                    type='range', value=self.value,
+                    min=self.min, max=self.max, step=self.step)
 
 
 class ColorPicker(Widget):
@@ -714,11 +702,8 @@ class ColorPicker(Widget):
         ''')
 
     def widget_html(self):
-        html_str = f'''<input {self.std_attrs()} type="color" value="{
-            self.value}">'''
-        if self.desc is not None:
-            html_str = f"""<span>{self.desc}</span>{html_str}"""
-        return html_str
+        return show.emit_tag('input', self.std_attrs(),
+                    type='color', value=self.value)
 
 
 class Choice(Widget):
@@ -726,13 +711,12 @@ class Choice(Widget):
     A set of radio button choices.
     """
 
-    def __init__(self, choices=None, selection=None, horizontal=False,
+    def __init__(self, choices=None, selection=None,
                  **kwargs):
         super().__init__(**kwargs)
         if choices is None:
             choices = []
         self.choices = Property(choices)
-        self.horizontal = Property(horizontal)
         self.selection = Property(selection)
 
     def widget_js(self):
@@ -748,7 +732,7 @@ class Choice(Widget):
               return '<label><input type="radio" name="choice" value="' +
                  esc(c) + '">' + esc(c) + '</label>'
             });
-            element.innerHTML = lines.join(model.get('horizontal')?' ':'<br>');
+            element.innerHTML = lines.join();
           }
           model.on('choices horizontal', render);
           model.on('selection', (ev) => {
@@ -762,14 +746,14 @@ class Choice(Widget):
         ''')
 
     def widget_html(self):
-        radios = [
-            f"""<label><input name="choice" type="radio" {
-            'checked' if value == self.selection else ''
-            } value="{html.escape(value)}">{html.escape(value)}</label>"""
-            for value in self.choices]
-        sep = " " if self.horizontal else "<br>"
-        return f'<form {self.std_attrs()}>{sep.join(radios)}</form>'
-
+        out = []
+        with show.enter_tag('form', self.std_attrs(), out=out):
+            for value in self.choices:
+                with show.enter_tag('label', out=out):
+                    show.emit_tag('input',
+                        (show.attrs(checked=None) if value == self.selection else None),
+                        name='choice', type='radio', value=value, out=out)
+        return ''.join(out)
 
 class Menu(Widget):
     """
@@ -798,7 +782,6 @@ class Menu(Widget):
             });
             element.menu.innerHTML = lines.join('\\n');
           }
-          model.on('choices horizontal', render);
           model.on('selection', (ev) => {
             [...element.querySelectorAll('option')].forEach((e) => {
               e.selected = (e.value == ev.value);
@@ -810,14 +793,15 @@ class Menu(Widget):
         ''')
 
     def widget_html(self):
-        options = [
-            f"""<option value="{html.escape(str(value))}" {
-            'selected' if value == self.selection else ''
-            }>{html.escape(str(value))}</option>"""
-            for value in self.choices]
-        sep = "\n"
-        return f'''<form {self.std_attrs()}"><select name="menu">{
-             sep.join(options)}</select></form>'''
+        out = []
+        with show.enter_tag('form', self.std_attrs(), out=out):
+            with show.enter_tag(show.Tag('select', name='menu'), out=out):
+                for value in self.choices:
+                    with show.enter_tag(show.Tag('option',
+                            (show.attrs(selected=None) if value == self.selection else None),
+                            value=value, out=out)):
+                        out.append(html.escape(str(value)))
+        return ''.join(out)
 
 
 class Datalist(Widget):
@@ -877,15 +861,14 @@ class Datalist(Widget):
         ''')
 
     def widget_html(self):
-        options = [
-            f"""<option value="{html.escape(str(value))}">"""
-            for value in self.choices]
-        return ''.join([
-            f'<form {self.std_attrs()} onsubmit="return false;">',
-            f'<input name="inp" list="{self.datalist_id()}" autocomplete="off">',
-            f'<datalist id="{self.datalist_id()}">',
-            ''.join(options),
-            f'</datalist></form>'])
+        out = []
+        with show.enter_tag('form', self.std_attrs(),
+                onsubmit='return false;', out=out):
+            show.emit_tag('input', name='inp', list=self.datalist_id(),
+                    autocomplete='off', out=out)
+            with show.enter_tag(show.Tag('datalist'), id=self.datalist_id()):
+                for value in self.choices:
+                    show.emit_tag('option', value=str(value))
 
 
 class Div(Widget):
@@ -933,7 +916,10 @@ class Div(Widget):
         ''')
 
     def widget_html(self):
-        return f'''<div {self.std_attrs()}>{self.innerHTML}</div>'''
+        out = []
+        with emit.enter_tag(self.std_attrs(), out=out):
+            out.append(self.innerHTML);
+        return ''.join(out)
 
 
 class ClickDiv(Div):
@@ -995,7 +981,8 @@ class Image(Widget):
         ''')
 
     def widget_html(self):
-        return f'''<img {self.std_attrs()} src="{html.escape(self.src)}">'''
+        return show.emit_tag('img', self.std_attrs(), show.style(margin=0),
+                src=self.src)
 
 
 ##########################################################################
diff --git a/baulab/paintwidget.py b/baulab/paintwidget.py
index fb129aa..08b4f3e 100644
--- a/baulab/paintwidget.py
+++ b/baulab/paintwidget.py
@@ -31,18 +31,18 @@ class PaintWidget(Widget):
 
     def widget_html(self):
         v = self.view_id()
-        return minify(f'''
-    <style>
-    #{v} {{ position: relative; display: inline-block; }}
-    #{v} .paintmask {{
-      position: absolute; top:0; left: 0; z-index: 1;
-      opacity: { self.opacity } }}
-    #{v} .paintmask.vanishing {{
-      opacity: 0; transition: opacity .1s ease-in-out; }}
-    #{v} .paintmask.vanishing:hover {{ opacity: { self.opacity }; }}
-    </style>
-    <div id="{v}"></div>
-    ''')
+        out = [minify(f'''
+            <style>
+            #{v} {{ position: relative; display: inline-block; }}
+            #{v} .paintmask {{
+              position: absolute; top:0; left: 0; z-index: 1;
+              opacity: { self.opacity } }}
+            #{v} .paintmask.vanishing {{
+              opacity: 0; transition: opacity .1s ease-in-out; }}
+            #{v} .paintmask.vanishing:hover {{ opacity: { self.opacity }; }}
+            </style>''')]
+        show.emit_tag(self.std_attrs(), out=out)
+        return ''.join(out)
 
 
 PAINT_WIDGET_JS = """
diff --git a/baulab/plotwidget.py b/baulab/plotwidget.py
index 67543b8..e4248c5 100644
--- a/baulab/plotwidget.py
+++ b/baulab/plotwidget.py
@@ -41,13 +41,13 @@ class PlotWidget(Image):
                 setattr(self, name, Property(default))
                 all_names.append(name)
 
-        old_backend = matplotlib.get_backend()
-        matplotlib.use('agg')
+        old_backend = matplotlib.pyplot.get_backend()
+        matplotlib.pyplot.switch_backend('agg')
         if 'mosaic' in init_args:
            self.fig, _ = matplotlib.pyplot.subplot_mosaic(**init_args)
         else:
            self.fig, _ = matplotlib.pyplot.subplots(**init_args)
-        matplotlib.use(old_backend)
+        matplotlib.pyplot.switch_backend(old_backend)
 
         def invoke_redraw():
             args = [self.fig]
diff --git a/baulab/show.py b/baulab/show.py
index 463865b..7621d9f 100644
--- a/baulab/show.py
+++ b/baulab/show.py
@@ -10,146 +10,393 @@
 
 import PIL.Image
 import base64
+import collections
+from contextlib import contextmanager
 import io
 import IPython
 import types
+import re
 import sys
-import html as html_module
-from IPython.display import display
+import inspect
+from html import escape
 
-g_buffer = None
+def show(*args):
+    '''
+    The main function.  Calls the IPython display function to show the
+    HTML-rendered arguments.
+    '''
+    display(html(*args))
 
+def html(*args):
+    '''
+    Renders the arguments into an HtmlString without displaying directly.
+    '''
+    out = []
+    with enter_tag(out=out):
+        for x in args:
+            render(x, out)
+    return HtmlString(''.join(out))
 
-def blocks(obj, space=''):
-    return IPython.display.HTML(space.join(blocks_tags(obj)))
+def bare_html(*args):
+    '''
+    Without an outermost element, renders the arguments as an HtmlString.
+    '''
+    out = []
+    for x in args:
+        render(x, out)
+    return HtmlString(''.join(out))
 
+def raw_html(*args):
+    '''
+    Produces an HtmlString from strings, without escaping or any fanciness.
+    '''
+    out = []
+    return HtmlString(''.join(str(x) for x in args))
 
-def rows(obj, space=''):
-    return IPython.display.HTML(space.join(rows_tags(obj)))
+@contextmanager
+def enter_tag(*args, out=None, **kwargs):
+    '''
+    Context manager for creating and emitting a tag and its matching
+    end-tag.  When the tag is created, any current defaults and options
+    to the styles and attributes are merged if present.
 
+    For example:
 
-def rows_tags(obj):
-    if isinstance(obj, dict):
-        obj = obj.items()
-    results = []
-    results.append('<table style="display:inline-table">')
-    for row in obj:
-        results.append('<tr style="padding:0">')
-        for item in row:
-            results.append('<td style="text-align:left; vertical-align:top;' +
-                           'padding:1px">')
-            results.extend(blocks_tags(item))
-            results.append('</td>')
-        results.append('</tr>')
-    results.append('</table>')
-    return results
+    ```
+    out = []
+    with show.enter_tag('div', id='d38', style(topMargin='8px'), out=out):
+        out.append('inside the div')
+    ```
 
+    The tags are merged with options, merged with the following precedence:
+        (1) Explicit tag, attr, or style options rentered before entering.
+        (2) Any tag, attr, and style specified on the `with` line.
+        (3) Child options specified by a parent.
 
-def blocks_tags(obj):
-    results = []
-    if hasattr(obj, '_repr_html_'):
-        results.append(obj._repr_html_())
-    elif isinstance(obj, PIL.Image.Image):
-        results.append(pil_to_html(obj))
-    elif isinstance(obj, (str, int, float)):
-        results.append('<div>')
-        results.append(html_module.escape(str(obj)))
-        results.append('</div>')
-    elif isinstance(obj, dict):
-        results.extend(blocks_tags([(k, v) for k, v in obj.items()]))
-    elif hasattr(obj, '__iter__'):
-        if hasattr(obj, 'tolist'):
-            # Handle numpy/pytorch tensors as lists.
-            try:
-                obj = obj.tolist()
-            except:
-                pass
-        blockstart, blockend, tstart, tend, rstart, rend, cstart, cend = [
-            '<div style="display:inline-block;text-align:center;line-height:1;' +
-            'vertical-align:top;padding:1px">',
-            '</div>',
-            '<table style="display:inline-table">',
-            '</table>',
-            '<tr style="padding:0">',
-            '</tr>',
-            '<td style="text-align:left; vertical-align:top; padding:1px">',
-            '</td>',
-        ]
-        needs_end = False
-        table_mode = False
-        for i, line in enumerate(obj):
-            if i == 0:
-                needs_end = True
-                if isinstance(line, tuple):
-                    table_mode = True
-                    results.append(tstart)
-                else:
-                    results.append(blockstart)
-            if table_mode:
-                results.append(rstart)
-                if not isinstance(line, str) and hasattr(line, '__iter__'):
-                    for cell in line:
-                        results.append(cstart)
-                        results.extend(blocks_tags(cell))
-                        results.append(cend)
-                else:
-                    results.append(cstart)
-                    results.extend(blocks_tags(line))
-                    results.append(cend)
-                results.append(rend)
+    Parent options are overwritten by `with` options, which can be
+    overwritten by a user who specifies explicit tag options when
+    rendering.
+    '''
+    global tag_stack
+    if len(tag_stack) and tag_stack[-1] is not None:
+        default_tag = tag_stack[-1]
+    else:
+        default_tag = HORIZONTAL
+    current_tag = default_tag(*args, **kwargs)
+    current_tag.update(*tag_modifications)
+    tag_modifications.clear()
+    tag_stack.append(current_tag.child)
+    try:
+        if out is not None:
+            out.append(current_tag.begin())
+        yield current_tag
+        if out is not None:
+            out.append(current_tag.end())
+    finally:
+        tag_modifications.clear()
+        tag_stack.pop()
+
+def emit_tag(*args, out=None, **kwargs):
+    '''
+    Emits the specified tag, applying any current defaults and options
+    to the styles and attributes.  Options are handled as with enter_tag.
+
+    If no `out` array is provided, returns the tag as a string.
+    '''
+    emit_out = [] if out is None else out
+    with enter_tag(*args, out=emit_out, **kwargs):
+        if out is None:
+            return ''.join(emit_out)
+
+HTML_EMPTY = set(('area base br col embed hr img input keygen '
+                  'link meta param source track wbr').split(' '))
+
+CSS_UNITS = dict([(k, unit) for keys, unit in [
+  ('width height min-width max-width min-height max-height', 'px'),
+  ('left right top bottom', 'px'),
+  ('font-size', 'px'),
+  ('border border-left border-right border-top border-bottom '
+    'border-width border-left-width border-right-width '
+    'border-top-width border-bottom-width', 'px'),
+  ('margin margin-left margin-right margin-top margin-bottom', 'px'),
+  ('padding padding-left padding-right padding-top padding-bottom', 'px'),
+] for k in keys.split(' ')])
+
+def hyphenateCamelKeys(d):
+    return {re.sub('([A-Z]+)', r'-\1', k).lower() : v for k, v in d.items()}
+
+def styleValue(v, k):
+    if isinstance(v, (int, float)) and k in CSS_UNITS:
+        return f'{v}{CSS_UNITS[k]}'
+    return str(v)
+
+class Style(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **hyphenateCamelKeys(kwargs))
+    def update(self, *args, **kwargs):
+        super().update(*args, **hyphenateCamelKeys(kwargs))
+    def __call__(self, **kwargs):
+        result = Style(**self, **hyphenateCamelKeys(kwargs))
+        return result
+    def __str__(self):
+        return ';'.join(f'{k}:{styleValue(v, k)}'
+                for k, v in self.items() if v is not None)
+
+def style(*args, **kwargs):
+    return Style(*args, **kwargs)
+
+class Attr(dict):
+    def __call__(self, **kwargs):
+        result = Attr(**self, **kwargs)
+        if 'style' in result and not result['style']:
+            del result['style']
+        return result
+    def __str__(self):
+        return ''.join(f' {k}' if v is None else f' {k}="{escape(str(v))}"'
+                for k, v in self.items())
+
+def attr(*args, **kwargs):
+    return Attr(*args, **kwargs)
+
+class ChildTag:
+    def __init__(self, child):
+        assert child is None or isinstance(child, Tag)
+        self.child = child
+
+class Tag:
+    def __init__(self, *args, **kwargs):
+        self.tag = 'div'
+        self.attrs = Attr()
+        self.style = Style()
+        self.child = None
+        self.update(*args, **kwargs)
+    def update(self, *args, **kwargs):
+        if len(args) > 0 and isinstance(args[0], str):
+            self.tag = args[0].lower()
+            args = args[1:]
+        for arg in args:
+            if isinstance(arg, Tag):
+                self.tag = arg.tag
+                self.attrs.update(arg.attrs)
+                self.style.update(arg.style)
+                self.child = arg.child
+            elif isinstance(arg, Attr):
+                self.attrs.update(arg)
+            elif isinstance(arg, Style):
+                self.style.update(arg)
+            elif isinstance(arg, ChildTag):
+                self.child = arg.child
+            elif arg is None:
+                continue
             else:
-                results.extend(blocks_tags(line))
-        if needs_end:
-            results.append(table_mode and tend or blockend)
-    return results
+                assert False, f'arg {arg} is not Tag or Attr or Style'
+        self.attrs.update(**kwargs)
+    def begin(self):
+        return f'<{self.tag}{self.attrs(style=str(self.style))}>'
+    def end(self):
+        return '' if self.tag in HTML_EMPTY else f'</{self.tag}>'
+    def __call__(self, *args, **kwargs):
+        result = Tag(self)
+        result.update(*args, **kwargs)
+        return result
+    def __str__(self):
+        return self.begin()
+    def __repr__(self):
+        return str(self)
 
 
-def pil_to_b64(img, format='png'):
-    buffered = io.BytesIO()
-    img.save(buffered, format=format)
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+# This is the default loop for nesting children: horizontal layout by default,
+# and then vertical layout for nested arrays; then horizontal within those, etc.
+HORIZONTAL = Tag(
+        style(display='flex', flex='1', flexFlow='row wrap', gap='3px',
+            alignItems='center'))
+VERTICAL = Tag(
+        style(display='flex', flex='1', flexFlow='column', gap='3px'),
+        ChildTag(HORIZONTAL))
+HORIZONTAL.update(ChildTag(VERTICAL))
+INLINE = Tag(
+        style(display='inline-flex', flexFlow='row wrap', gap='3px',
+            alignItems='center'),
+        ChildTag(VERTICAL))
+
+tag_stack = [INLINE]
+tag_modifications = []
+
+def modify_tag(*args):
+    tag_modifications.extend(args)
+
+def render(obj, out):
+    for detector, renderer in RENDERING_RULES:
+        if (isinstance(obj, detector) if isinstance(detector, (type, tuple)) else detector(obj)):
+            if renderer(obj, out) is not False:
+                return
+    # Fallback: convert object to string and then apply HTML escaping.
+    render_str(obj, out)
+
+def render_str(obj, out):
+    '''
+    Strings must be escaped.
+    '''
+    s = str(obj)
+    if '\n' in s:
+        render_pre(s, out)
+        return
+    with enter_tag(out=out):
+        out.append(escape(s))
+
+def render_html(obj, out):
+    '''
+    Use _repr_html_() when available and non-None.
+    '''
+    h = obj._repr_html_()
+    if h is None:
+        return False
+    out.append(h)
+
+def render_list(obj, out):
+    '''
+    Lists are divs containin divs, alternating row-inline and column flex layout.
+    '''
+    with enter_tag(out=out):
+        for v in obj:
+            render(v, out)
+
+def render_dict(obj, out):
+    '''
+    Dicts become tables.
+    '''
+    with enter_tag('table', style(display=None, width='100%'), ChildTag(None), out=out):
+        for k, v in obj.items():
+            with enter_tag('tr', out=out):
+                out.append(row.begin())
+                with enter_tag('td', ChildTag(HORIZONTAL), out=out):
+                    out.append(escape(str(k)))
+                with enter_tag('td', ChildTag(HORIZONTAL), out=out):
+                    render(v, out)
+
+def render_image(obj, out):
+    '''
+    Images and figures become <img> tags.
+    '''
+    try:
+        buf = io.BytesIO()
+        if hasattr(obj, 'save'): # Like a PIL.Image.Image
+            obj.save(buf, format='png')
+        elif hasattr(obj, 'savefig'): # Like a matplotlib.figure.Figure
+            obj.savefig(buf, format='png', bbox_inches='tight')
+        else:
+            assert False
+        src = 'data:image/png;base64,' + (
+                base64.b64encode(buf.getvalue()).decode('utf-8'))
+        buf.close()
+    except:
+        return False
+    emit_tag('img', attr(src=src), style(flex='0'), out=out)
+
+def render_pre(obj, out):
+    '''
+    Long multiline text data types are rendered in <pre> tags.
+    '''
+    s = str(obj)
+    with enter_tag('pre', out=out):
+        out.append(escape(s))
+
+def render_modifications(obj, out):
+    '''
+    Allows Tag, Attr, Style, ChildTag objects to modify the next tag to output.
+    '''
+    assert isinstance(obj, (Tag, Attr, Style, ChildTag))
+    modify_tag(obj)
+
+def subclass_of(clsname):
+    '''
+    Detects if obj is a subclass of a class named clsname, without requiring import
+    of the class.
+    '''
+    def test(obj):
+        for x in inspect.getmro(type(obj)):
+            if clsname == x.__module__ + '.' + x.__name__:
+                return True
+        return False
+    return test
+
+RENDERING_RULES = [
+        # Special tag modifications
+        ((Style, Attr, Tag, ChildTag), render_modifications),
+        # Objects with an html repr
+        ((lambda x: hasattr(x, '_repr_html_')), render_html),
+        # Strings should not be treated as lists
+        (str, render_str),
+        # Dictionaries: render as table
+        (collections.abc.Mapping, render_dict),
+        # PIL images
+        (subclass_of('PIL.Image.Image'), render_image),
+        # Matplotlib figures
+        (subclass_of('matplotlib.figure.Figure'), render_image),
+        # Numpy, pytorch arrays
+        (lambda x: hasattr(x, 'shape') and hasattr(x, 'dtype'), render_pre),
+        # Generators and lists: recurse
+        ((lambda x: hasattr(x, '__iter__')), render_list),
+]
 
 
-def pil_to_url(img, format='png'):
-    return 'data:image/%s;base64,%s' % (format, pil_to_b64(img, format))
-
-
-def pil_to_html(img, margin=1):
-    mattr = ' style="margin:%dpx"' % margin
-    return '<img src="%s"%s>' % (pil_to_url(img), mattr)
-
-
-def a(x, cols=None):
-    global g_buffer
-    if g_buffer is None:
-        g_buffer = []
-    g_buffer.append(x)
-    if cols is not None and len(g_buffer) >= cols:
-        flush()
-
-
-def reset():
-    global g_buffer
-    g_buffer = None
-
-
-def flush(*args, **kwargs):
-    global g_buffer
-    if g_buffer is not None:
-        x = g_buffer
-        g_buffer = None
-        display(blocks(x, *args, **kwargs))
-
-
-def show(x=None, *args, **kwargs):
-    flush(*args, **kwargs)
-    if x is not None:
-        display(blocks(x, *args, **kwargs))
-
-
-def html(obj, space=''):
-    return blocks(obj, space)._repr_html_()
-
+class HtmlString(str):
+    '''
+    A string that contains HTML, and that returns itself as _repr_html_.
+    It does no escaping, and just interprets strings as markup.
+    '''
+    def __new__(cls, s):
+        return super().__new__(cls, s)
+    def _repr_html_(self):
+        return self
+    def __add__(self, other):
+        return HtmlString(super().__add__(other))
+    def __mul__(self, other):
+        return HtmlString(super().__mul__(other))
+    def __rmul__(self, other):
+        return HtmlString(super().__rmul__(other))
+    def __mod__(self, other):
+        return HtmlString(super().__mod__(other))
+    def format(self, *args, **kwargs):
+        return HtmlString(super().format(*args, **kwargs))
+    def format_map(self, mapping):
+        return HtmlString(super().format_map(mapping))
+    def expandtabs(self, tabsize=8):
+        return HtmlString(super().expandtabs(tabsize=tabsize))
+    def join(self, iterable):
+        return HtmlString(super().join(iterable))
+    def ljust(self, width, fillchar=' '):
+        return HtmlString(super().ljust(width, fillchar=fillchar))
+    def lstrip(self, chars=None):
+        return HtmlString(super().lstrip(chars))
+    def removeprefix(self, prefix):
+        return HtmlString(super().removeprefix(prefix))
+    def removesuffix(self, suffix):
+        return HtmlString(super().removesuffix(suffix))
+    def replace(self, old, new, count=-1):
+        return HtmlString(super().replace(old, new, count=count))
+    def rjust(self, width, fillchar=' '):
+        return HtmlString(super().rjust(width, fillchar=fillchar))
+    def rstrip(self, chars=None):
+        return HtmlString(super().rstrip(chars))
+    def strip(self, chars=None):
+        return HtmlString(super().strip(chars))
+    def translate(self, table):
+        return HtmlString(super().translate(table))
+    def capitalize(self):
+        return HtmlString(super().capitalize())
+    def casefold(self):
+        return HtmlString(super().casefold())
+    def swapcase(self):
+        return HtmlString(super().swapcase())
+    def lower(self):
+        return HtmlString(super().lower())
+    def title(self):
+        return HtmlString(super().title())
+    def upper(self):
+        return HtmlString(super().upper())
+    def zfill(self, width):
+        return HtmlString(super().zfill(width))
 
 class CallableModule(types.ModuleType):
     def __init__(self):