Match the surrounding page's theme automatically

docs/api/treescope.rst

@@ -105,6 +105,7 @@ configured globally using `context.ContextualValue.set_globally`.
   default_diverging_colormap
   default_sequential_colormap
   active_autovisualizer
+  theme
   active_renderer
   default_magic_autovisualizer
   active_expansion_strategy

tests/representation_test.py

@@ -23,6 +23,7 @@
 import jax
 from treescope import lowering
 from treescope import type_registries
+from treescope._internal.api import theme
 from treescope._internal.parts import basic_parts
 from treescope._internal.parts import common_structures
 from treescope._internal.parts import part_interface
@@ -1187,6 +1188,53 @@ def test_render_as_root(self):
           ['[mock A, b:y]', '[mock B, b:y]'], html_output
       )
 
+  def test_theme(self):
+    part = MockRenderableTreePart('A')
+    with self.subTest('light_by_default'):
+      self.assertIn(
+          'class="treescope_root theme_light"',
+          lowering.render_to_html_as_root(part),
+      )
+    with self.subTest('auto'):
+      with theme.theme.set_scoped('auto'):
+        self.assertIn(
+            'class="treescope_root theme_auto"',
+            lowering.render_to_html_as_root(part),
+        )
+    with self.subTest('light'):
+      with theme.theme.set_scoped('light'):
+        self.assertIn(
+            'class="treescope_root theme_light"',
+            lowering.render_to_html_as_root(part),
+        )
+    with self.subTest('dark'):
+      with theme.theme.set_scoped('dark'):
+        self.assertIn(
+            'class="treescope_root theme_dark"',
+            lowering.render_to_html_as_root(part),
+        )
+    with self.subTest('roundtrip_mode_is_preserved'):
+      with theme.theme.set_scoped('dark'):
+        self.assertIn(
+            'class="treescope_root theme_dark roundtrip_mode"',
+            lowering.render_to_html_as_root(part, roundtrip=True),
+        )
+    with self.subTest('invalid'):
+      with theme.theme.set_scoped('sepia'):
+        with self.assertRaisesRegex(ValueError, "but got 'sepia'"):
+          lowering.render_to_html_as_root(part)
+    with self.subTest('auto_detection_hook'):
+      # The container sets a dataset property and the stylesheet keys off the
+      # corresponding attribute, so the two have to keep agreeing with each
+      # other even though they are spelled differently.
+      html_output = lowering.render_to_html_as_root(part)
+      self.assertIn(
+          'root.dataset.treescopeDetectedTheme = pageIsDark ? "dark" : "light"',
+          html_output,
+      )
+      self.assertIn('[data-treescope-detected-theme="dark"]', html_output)
+      self.assertIn('[data-treescope-detected-theme]', html_output)
+
   def test_build_qualified_type_name(self):
     type_registries.update_registries_for_imports()
     part = common_structures.maybe_qualified_type_name(jax.ShapeDtypeStruct)

treescope/__init__.py

@@ -70,6 +70,9 @@
     Autovisualizer,
     active_autovisualizer,
 )
+from ._internal.api.theme import (
+    theme,
+)
 from ._internal.api.default_renderer import (
     active_renderer,
     active_expansion_strategy,
@@ -121,6 +124,27 @@
 choice; a common choice is arrayviz's `ArrayAutovisualizer()`.
 """
 
+theme: context.ContextualValue[typing.Literal['light', 'dark', 'auto']] = theme
+"""The theme to use when rendering a tree to HTML.
+
+This can be set to:
+
+* ``"light"`` (the default), to render dark text on a light background,
+* ``"dark"``, to render light text on a dark background,
+* ``"auto"``, to match the theme of whatever is displaying the
+  rendering.
+
+In ``"auto"`` mode, the theme is determined in the browser when the rendering
+is displayed, by inspecting the text color that the rendering inherits from
+the page around it. Renderings fall back to ``"light"`` if that isn't
+possible, for instance if they are displayed somewhere that doesn't run
+JavaScript. A rendering that has matched the page also keeps that page's
+background instead of drawing one of its own, so that it blends in with
+whatever shade the page happens to use.
+
+This has no effect on text renderings, which are never colored.
+"""
+
 active_renderer: context.ContextualValue[renderers.TreescopeRenderer] = (
     active_renderer
 )

treescope/_internal/api/theme.py

@@ -0,0 +1,28 @@
+# Copyright 2024 The Treescope Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Configurable settings for the theme of renderings."""
+
+from typing import Literal
+
+from treescope import context
+
+
+Theme = Literal["light", "dark", "auto"]
+
+theme: context.ContextualValue[Theme] = context.ContextualValue(
+    module=__name__,
+    qualname="theme",
+    initial_value="light",
+)

treescope/_internal/arrayviz_impl.py

@@ -79,7 +79,7 @@ def html_setup() -> (
         }
         .arrayviz_container .info input[type="range"] {
             vertical-align: middle;
-            filter: grayscale(1) opacity(0.5);
+            filter: grayscale(1) opacity(var(--treescope-slider-opacity));
         }
         .arrayviz_container .info input[type="range"]:hover {
             filter: grayscale(0.5);

treescope/_internal/html_encapsulation.py

@@ -141,6 +141,30 @@ class RunHere extends HTMLElement {
   .filter((elt) => !elt.dataset.setup)
 )[0];
 root.dataset.setup = 1;
+// Detect which theme the page around this rendering uses, so that renderings
+// using the "auto" theme can match it. We can't use the
+// `prefers-color-scheme` media query, because notebooks usually have a theme
+// setting of their own that is independent of the browser's, and because some
+// of them report the query incorrectly inside their output cells. Instead we
+// look at the text color we inherit from the page: if the surrounding text is
+// light, it was meant to be read against a dark background. Colors are
+// inherited across the shadow root boundary, so this has to run before the
+// rendering sets a color of its own. The color can be in any CSS color format
+// (e.g. `oklch(...)`), so we draw it onto a canvas and read back the pixel to
+// let the browser convert it to RGB. If we can't tell, we leave the attribute
+// unset, and the rendering falls back to its own light background.
+const inheritedColor = window.getComputedStyle(root).color;
+const colorCtx = document.createElement("canvas").getContext("2d");
+if (colorCtx) {
+  colorCtx.fillStyle = inheritedColor;
+  colorCtx.fillRect(0, 0, 1, 1);
+  const [r, g, b] = colorCtx.getImageData(0, 0, 1, 1).data;
+  // Perceptual luminance, as in arrayviz's contrasting-color logic.
+  const pageIsDark = 0.2126 * r + 0.7152 * g + 0.0722 * b > 128;
+  root.dataset.treescopeDetectedTheme = pageIsDark ? "dark" : "light";
+} else {
+  console.warn("Treescope: could not detect the page theme.");
+}
 const msg = document.createElement("span");
 msg.style = "color: #cccccc; font-family: monospace;";
 msg.textContent = "(Loading...)";

treescope/lowering.py

@@ -34,6 +34,7 @@
 from treescope._internal.parts import foldable_impl
 from treescope._internal.parts import part_interface
 from treescope._internal.api import abbreviation
+from treescope._internal.api import theme as theme_lib
 
 _deferrables: context.ContextualValue[
     list[foldable_impl.DeferredWithThunk] | None
@@ -328,6 +329,7 @@ def _render_one(
       --treescope-outline-color: #aaaaaa;
       --treescope-axis-label-color: black;
       --treescope-info-color: #aaaaaa;
+      --treescope-slider-opacity: 0.5;
       --treescope-tooltip-bg: white;
       --treescope-tooltip-fg: black;
       --treescope-tooltip-border-color: black;
@@ -347,6 +349,35 @@ def _render_one(
       content-visibility: auto;
       contain-intrinsic-size: auto none;
     }
+    .treescope_root.theme_dark,
+    :host([data-treescope-detected-theme="dark"]) .treescope_root.theme_auto {
+      color-scheme: dark;
+      --treescope-bg: #1f1f1f;
+      --treescope-fg: #cccccc;
+      --treescope-keyword-color: #569cd6;
+      --treescope-number-color: #b5cea8;
+      --treescope-string-color: #ce9178;
+      --treescope-comment-color: #a6a6a6;
+      --treescope-abbreviation-color: var(--treescope-string-color);
+      --treescope-error-color: #f85149;
+      --treescope-deferred-color: #a6a6a6;
+      --treescope-marker-color: #6e7681;
+      --treescope-faint-color: #3f3f3f;
+      --treescope-copy-button-color: #6a6a6a;
+      --treescope-outline-color: #6a6a6a;
+      --treescope-axis-label-color: #cccccc;
+      --treescope-info-color: #9d9d9d;
+      --treescope-slider-opacity: 0.75;
+      --treescope-tooltip-bg: #202020;
+      --treescope-tooltip-fg: #cccccc;
+      --treescope-tooltip-border-color: #313131;
+      --treescope-input-color: #989898;
+    }
+    :host([data-treescope-detected-theme]) .treescope_root.theme_auto {
+      /* Keep the page's own background, since themes disagree about which
+         off-white or dark gray to use. */
+      --treescope-bg: transparent;
+    }
   """))
   stream.write("</style>")
   # These scripts allow us to defer execution of javascript blocks until after
@@ -358,7 +389,13 @@ def _render_one(
   stream.write("</script></treescope-run-here>")
 
   # Render the root node.
-  classnames = "treescope_root"
+  requested_theme = theme_lib.theme.get()
+  if requested_theme not in ("light", "dark", "auto"):
+    raise ValueError(
+        "`treescope.theme` must be 'light', 'dark', or 'auto', but got"
+        f" {requested_theme!r}"
+    )
+  classnames = f"treescope_root theme_{requested_theme}"
   if roundtrip:
     classnames += " roundtrip_mode"
   stream.write(