Defining `_repr_html_` for Pytree dataclasses which uses `pz.ts.render_to_html` leads to infinite loop

[femtomc]

Hi! I'm attempting to setup pretty printing for some classes in my library using treescope for notebooks / IPython.

I defined some custom handlers for my classes, and overloaded the active renderer with my handler. Now, I want to define something like this:

def _repr_html_(self):
   ...
   return pz.ts.render_to_html(self)

after overloading the active renderer, so Penzai will use my custom handlers, and return HTML. The problem is that render_to_html will call _repr_html_ (because I'm defining it for my class, and render_to_html checks if this exists as an attribute, and then calls it), leading to an infinite loop.

Is there a recommended pattern for this?

[danieldjohnson]

Ah, that's an interesting interaction. Right now Treescope assumes things that define _repr_html_ are "figure-like", so render_to_html computes and adds the rich visualization as an annotation. This is useful for stuff like showing lists of plotly figures. I didn't anticipate that someone would call render_to_html inside _repr_html_.

This behavior is enabled by the repr_html_postprocessor.append_repr_html_when_present wrapper hook (here). You can disable this behavior by removing that wrapper hook from the wrapper_hooks list.

In this case, you might also want to avoid using pz.ts.render_to_html directly, because that uses the default global renderer which you might not always have control over. Instead, you could construct a new renderer specific to your types, and call renderer.to_html instead. That would let you fully control the set of handlers that are enabled for the renderer you're using inside _repr_html_.

(For context, are you doing this so that you can render your objects with treescope even if treescope isn't enabled by default? If treescope is enabled, you shouldn't have to define _repr_html_ at all.)

[femtomc]

@danieldjohnson Got it -- I think I'll just install treescope by default in my notebooks, and not futz around with _repr_html_ at all.

In that case, is there a pithy few line incantation for me to install treescope and push my own custom handler onto the default?

[danieldjohnson]

Sounds good, if that works for you then that's the approach I'd recommend. I'll go ahead and close this issue then.

To register treescope as the default pretty printer and also change the handler, you could run something like this:

# Install things in IPython
pz.ts.register_as_default()
pz.ts.register_autovisualize_magic()

# Enable global overrides
pz.enable_interactive_context()

# Configure the custom handler
my_custom_renderer = pz.ts.active_renderer.get().extended_with(<stuff here>)
pz.ts.active_renderer.set_interactive(my_custom_renderer)

# Optional: enable automatic array visualization
pz.ts.active_autovisualizer.set_interactive(pz.ts.ArrayAutovisualizer())

You could alternatively set my_custom_renderer to a custom instance of penzai.treescope.renderer.TreescopeRenderer if you want to fully control its configuration.