pyglove 0.4.5.dev202410170809__py3-none-any.whl → 0.4.5.dev202410190807__py3-none-any.whl

pyglove/core/views/html/tree_view.py

@@ -14,8 +14,7 @@
 """HTML Tree View (The default view for PyGlove objects)."""
 
 import inspect
-import re
-from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Union
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Sequence, Tuple, Union
 
 from pyglove.core import object_utils
 from pyglove.core.views.html import base
@@ -38,93 +37,84 @@ class HtmlTreeView(HtmlView):
   class Extension(HtmlView.Extension):
     """The base class for extensions for HtmlTreeView."""
 
-    #
-    # Default extension-level rendering options overrides.
-    #
+    def _html_tree_view_render(
+        self,
+        *,
+        view: 'HtmlTreeView',
+        name: Optional[str] = None,
+        parent: Any = None,
+        root_path: Optional[KeyPath] = None,
+        **kwargs,
+    ) -> Html:
+      """The entrypoint of rendering the subtree represented by this extension.
 
-    def _html_tree_view_special_keys(self) -> Sequence[str]:
-      """Returns the special keys to display (at the immediate child level)."""
-      return []
+      Args:
+        view: The view to render the object.
+        name: The name of the object.
+        parent: The parent of the object.
+        root_path: The key path of the object relative to the root.
+        **kwargs: kwargs to pass to `view.render()` on this extension.
 
-    def _html_tree_view_include_keys(self) -> Optional[Sequence[str]]:
-      """Returns the keys to include (at the immediate child level)."""
-      return None
+      Returns:
+        The rendered HTML.
+      """
+      return self._html_tree_view(
+          view=view,
+          name=name,
+          parent=parent,
+          root_path=root_path,
+          **view.get_kwargs(
+              kwargs, self._html_tree_view_config(), root_path or KeyPath()
+          )
+      ).add_style(
+          *self._html_tree_view_css_styles()
+      )
 
-    def _html_tree_view_exclude_keys(self) -> Sequence[str]:
-      """Returns the keys to include (at the immediate child level)."""
+    #
+    # Users could override this methods to customize the styles and
+    # rendering arguments for the subtree.
+    #
+
+    @classmethod
+    def _html_tree_view_css_styles(cls) -> list[str]:
+      """Returns the CSS styles for the subtree."""
+      del cls
       return []
 
-    def _html_tree_view_uncollapse_level(self) -> Optional[int]:
-      """Returns the level of the subtree to uncollapse.
+    @classmethod
+    def _html_tree_view_config(cls) -> Dict[str, Any]:
+      """Returns the config (rendering arguments) of current extension.
 
       Returns:
-        The level of subtree to uncollapse. If None, the subtree will be fully
-        expanded. Please note that the uncollapsed subtree will show only when
-        current node is uncollapsed.
+        A dictionary of rendering arguments for the subtree. These arguments
+        will override the arguments passed to `view.render()`. See the
+        `render()` method for the full list of arguments.
       """
-      return 1
-
-    def _html_tree_view_uncollapse(self) -> KeyPathSet:
-      """Returns the node paths (relative to current node) to uncollapse."""
-      return KeyPathSet()
+      return {}
 
     #
-    # Default behavior overrides.
+    # Users could override the methods below to customize rendering
+    # logics.
     #
 
-    def _html_tree_view_render(
+    def _html_tree_view(
         self,
         *,
-        view: 'HtmlView',
-        name: Optional[str],
-        parent: Any,
-        root_path: KeyPath,
-        title: Union[str, Html, None] = None,
-        special_keys: Optional[Sequence[Union[int, str]]] = None,
-        include_keys: Optional[Iterable[Union[int, str]]] = None,
-        exclude_keys: Optional[Iterable[Union[int, str]]] = None,
-        collapse_level: Optional[int] = HtmlView.PresetArgValue(1),
-        uncollapse: Union[
-            KeyPathSet, base.NodeFilter, None
-        ] = HtmlView.PresetArgValue(None),
-        filter: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),  # pylint: disable=redefined-builtin
-        highlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-        lowlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-        enable_summary: Optional[bool] = HtmlView.PresetArgValue(None),
-        max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-        enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-        enable_key_tooltip: bool = HtmlView.PresetArgValue(True),
+        view: 'HtmlTreeView',
+        name: Optional[str] = None,
+        parent: Any = None,
+        root_path: Optional[KeyPath] = None,
         **kwargs,
     ) -> Html:
-      """Returns the topmost HTML representation of the object.
+      """Returns the topmost HTML representation of this extension.
 
       Args:
         view: The view to render the object.
         name: The name of the object.
         parent: The parent of the object.
         root_path: The key path of the object relative to the root.
-        title: The title of the summary.
-        special_keys: The special keys to display (at the immediate child
-          level).
-        include_keys: The keys to include (at the immediate child level).
-        exclude_keys: The keys to exclude (at the immediate child level).
-        collapse_level: The level to collapse the tree (relative to this node).
-        uncollapse: A key path set (relative to root_path) for the nodes to
-          uncollapse. or a function with signature (path, value, parent) -> bool
-          to filter nodes to uncollapse.
-        filter: A function with signature (path, value, parent) -> include
-          to determine whether to include a field (at all levels).
-        highlight: A function with signature (path, value, parent) -> bool
-          to determine whether to highlight a field.
-        lowlight: A function with signature (path, value, parent) -> bool
-          to determine whether to lowlight a field.
-        enable_summary: Whether to enable the summary. If None, summary will
-          be enabled for complex types or when string exceeds
-          `max_summary_len_for_str`.
-        max_summary_len_for_str: The maximum length of the string to display.
-        enable_summary_tooltip: Whether to enable the tooltip for the summary.
-        enable_key_tooltip: Whether to enable the tooltip for the key.
-        **kwargs: Additional keyword arguments passed from `pg.to_html`.
+        **kwargs: kwargs to pass to the view. See `_html_tree_view_config` for
+          the builtin arguments.
 
       Returns:
         The rendered HTML.
@@ -134,65 +124,38 @@ class HtmlTreeView(HtmlView):
           name=name,
           parent=parent,
           root_path=root_path,
-          title=title,
-          special_keys=special_keys,
-          include_keys=include_keys,
-          exclude_keys=exclude_keys,
-          filter=filter,
-          highlight=highlight,
-          lowlight=lowlight,
-          enable_summary=enable_summary,
-          max_summary_len_for_str=max_summary_len_for_str,
-          enable_summary_tooltip=enable_summary_tooltip,
-          enable_key_tooltip=enable_key_tooltip,
-          collapse_level=collapse_level,
-          uncollapse=uncollapse,
-          **kwargs
+          **kwargs,
       )
 
     def _html_tree_view_summary(
         self,
         *,
         view: 'HtmlTreeView',
-        name: Optional[str],
-        parent: Any,
-        root_path: KeyPath,
-        title: Union[str, Html, None] = None,
-        enable_summary: Optional[bool] = HtmlView.PresetArgValue(None),
-        max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-        enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
+        name: Optional[str] = None,
+        parent: Any = None,
+        root_path: Optional[KeyPath] = None,
         **kwargs,
     ) -> Optional[Html]:
-      """Returns the HTML representation of the object.
+      """Returns the HTML summary for the object.
 
       Args:
         view: The view to render the object.
         name: The name of the object.
         parent: The parent of the object.
         root_path: The key path of the object relative to the root.
-        title: The title of the summary.
-        enable_summary: Whether to enable the summary. If None, summary will
-          be enabled for complex types or when string exceeds
-          `max_summary_len_for_str`.
-        max_summary_len_for_str: The maximum length of the string to display.
-        enable_summary_tooltip: Whether to enable the tooltip for the summary.
-        **kwargs: Additional keyword arguments passed from `pg.to_html`. These
-          arguments may be handled by the user logic but not the general
-          HtmlTreeView.
+        **kwargs: kwargs to pass to the view. See `_html_tree_view_config` for
+          the builtin arguments.
 
       Returns:
         An optional HTML object representing the summary of the object. If None,
-        the summary will be hidden.
+        the content will be returned directly instead of having a <details>
+        container.
       """
       return view.summary(
           self,
           name=name,
           parent=parent,
           root_path=root_path,
-          title=title,
-          enable_summary=enable_summary,
-          max_summary_len_for_str=max_summary_len_for_str,
-          enable_summary_tooltip=enable_summary_tooltip,
           **kwargs,
       )
 
@@ -200,22 +163,9 @@ class HtmlTreeView(HtmlView):
         self,
         *,
         view: 'HtmlTreeView',
-        name: Optional[str],
-        parent: Any,
-        root_path: KeyPath,
-        special_keys: Optional[Sequence[Union[int, str]]] = None,
-        include_keys: Optional[Iterable[Union[int, str]]] = None,
-        exclude_keys: Optional[Iterable[Union[int, str]]] = None,
-        collapse_level: Optional[int] = HtmlView.PresetArgValue(1),
-        uncollapse: Union[
-            KeyPathSet, base.NodeFilter, None
-        ] = HtmlView.PresetArgValue(None),
-        filter: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),  # pylint: disable=redefined-builtin
-        highlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-        lowlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-        max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-        enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-        enable_key_tooltip: bool = HtmlView.PresetArgValue(True),
+        name: Optional[str] = None,
+        parent: Any = None,
+        root_path: Optional[KeyPath] = None,
         **kwargs,
         ) -> Html:
       """Returns the main content for the object.
@@ -225,26 +175,8 @@ class HtmlTreeView(HtmlView):
         name: The name of the object.
         parent: The parent of the object.
         root_path: The key path of the object relative to the root.
-        special_keys: The special keys to display (at the immediate child
-          level).
-        include_keys: The keys to include (at the immediate child level).
-        exclude_keys: The keys to exclude (at the immediate child level).
-        collapse_level: The level to collapse the tree (relative to this node).
-        uncollapse: A key path set (relative to root_path) for the nodes to
-          uncollapse. or a function with signature (path, value, parent) -> bool
-          to filter nodes to uncollapse.
-        filter: A function with signature (path, value, parent) -> include
-          to determine whether to include a field (at all levels).
-        highlight: A function with signature (path, value, parent) -> bool
-          to determine whether to highlight a field (at all levels).
-        lowlight: A function with signature (path, value, parent) -> bool
-          to determine whether to lowlight a field (at all levels).
-        max_summary_len_for_str: The maximum length of the string to display.
-        enable_summary_tooltip: Whether to enable the tooltip for the summary.
-        enable_key_tooltip: Whether to enable the key tooltip.
-        **kwargs: Additional keyword arguments passed from `pg.to_html`. These
-          arguments may be handled by the user logic but not the general
-          HtmlTreeView.
+        **kwargs: kwargs to pass to the view. See `_html_tree_view_config` for
+          the builtin arguments.
 
       Returns:
         The rendered HTML as the main content of the object.
@@ -254,51 +186,11 @@ class HtmlTreeView(HtmlView):
           name=name,
           parent=parent,
           root_path=root_path,
-          special_keys=special_keys,
-          include_keys=include_keys,
-          exclude_keys=exclude_keys,
-          filter=filter,
-          highlight=highlight,
-          lowlight=lowlight,
-          max_summary_len_for_str=max_summary_len_for_str,
-          enable_summary_tooltip=enable_summary_tooltip,
-          enable_key_tooltip=enable_key_tooltip,
-          collapse_level=collapse_level,
-          uncollapse=uncollapse,
-          **kwargs
-      )
-
-    def _html_tree_view_tooltip(
-        self,
-        *,
-        view: 'HtmlTreeView',
-        name: Optional[str],
-        parent: Any,
-        root_path: KeyPath,
-        content: Union[str, Html, None] = None,
-        **kwargs,
-    ) -> Optional[Html]:
-      """Returns the tooltip for the object.
-
-      Args:
-        view: The view to render the object.
-        name: The referenced name of the object.
-        parent: The parent of the object.
-        root_path: The key path of the object relative to the root.
-        content: Custom content to display in the tooltip.
-        **kwargs: Additional keyword arguments passed from `pg.to_html`. These
-          arguments may be handled by the user logic but not the general
-          HtmlTreeView.
-
-      Returns:
-        An optional HTML object representing the tooltip of the object. If None,
-        the tooltip will be hidden.
-      """
-      return view.tooltip(
-          value=self, name=name, parent=parent, root_path=root_path,
-          content=content, **kwargs
+          **kwargs,
       )
 
+  # NOTE(daiyip): update `get_kwargs()` and `get_passthrough_kwargs()` when new
+  # arguments are added.
   @HtmlView.extension_method('_html_tree_view_render')
   def render(
       self,
@@ -307,22 +199,47 @@ class HtmlTreeView(HtmlView):
       name: Optional[str] = None,
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
+      css_classes: Optional[Sequence[str]] = None,
+      # Summary settings.
       title: Union[str, Html, None] = None,
-      special_keys: Optional[Sequence[Union[int, str]]] = None,
-      include_keys: Optional[Iterable[Union[int, str]]] = None,
-      exclude_keys: Optional[Iterable[Union[int, str]]] = None,
-      collapse_level: Optional[int] = HtmlView.PresetArgValue(1),
-      uncollapse: Union[
-          KeyPathSet, base.NodeFilter, None
-      ] = HtmlView.PresetArgValue(None),
-      filter: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),  # pylint: disable=redefined-builtin
-      highlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      lowlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      enable_summary: Optional[bool] = HtmlView.PresetArgValue(None),
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-      enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-      enable_key_tooltip: bool = HtmlView.PresetArgValue(True),
-      **kwargs
+      enable_summary: Optional[bool] = None,
+      enable_summary_for_str: bool = True,
+      max_summary_len_for_str: int = 80,
+      enable_summary_tooltip: bool = True,
+      summary_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      # Content settings.
+      key_style: Union[
+          Literal['label', 'summary'],
+          Callable[[KeyPath, Any, Any], Literal['label', 'summary']]
+      ] = 'summary',
+      key_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      include_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      exclude_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      enable_key_tooltip: bool = True,
+      # Collapse settings.
+      collapse_level: Optional[int] = 1,
+      uncollapse: Union[KeyPathSet, base.NodeFilter, None] = None,
+      # Extension settings.
+      child_config: Optional[Dict[str, Any]] = None,
+      extra_flags: Optional[Dict[str, Any]] = None,
+      # Tree operations.
+      highlight: Optional[base.NodeFilter] = None,
+      lowlight: Optional[base.NodeFilter] = None,
+      debug: bool = False,
   ) -> Html:
     """Renders the entire HTML tree view for the value.
 
@@ -331,102 +248,177 @@ class HtmlTreeView(HtmlView):
       name: The name of the value.
       parent: The parent of the value.
       root_path: The root path of the value.
-      title: The title of the summary.
-      special_keys: The special keys to display (at the immediate child level).
-      include_keys: The keys to include (at the immediate child level).
-      exclude_keys: The keys to exclude (at the immediate child level).
-      collapse_level: The level to collapse the tree (relative to this node).
-      uncollapse: A key path set (relative to root_path) for the nodes to
-        uncollapse. or a function with signature (path, value, parent) -> bool
-        to filter nodes to uncollapse.
-      filter: A function with signature (path, value, parent) -> include
-        to determine whether to include a field (at all levels).
-      highlight: A function with signature (path, value, parent) -> bool
-        to determine whether to highlight a field.
-      lowlight: A function with signature (path, value, parent) -> bool
-        to determine whether to lowlight a field.
+      css_classes: CSS classes to add to the top-most element.
+      title: The title of the summary. If None, the default title will be used,
+        which is the type name of the value.
       enable_summary: Whether to enable the summary. If None, summary will
         be enabled for complex types or when string exceeds
         `max_summary_len_for_str`.
+      enable_summary_for_str: Whether to enable the summary for strings.
       max_summary_len_for_str: The maximum length of the string to display.
       enable_summary_tooltip: Whether to enable the tooltip for the summary.
-      enable_key_tooltip: Whether to enable the key tooltip.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+      summary_color: The color used for the summary for displaying the referred
+        field name of the object. It can be a tuple of (color, background-color)
+        or a function that takes (root_path, value, parent) and returns the
+        color tuple.
+      key_style: The style of the key. If 'label', the key will be rendered as a
+        label. If 'summary', it will be rendered as a summary in the <details>
+        tag. If a function, it will be called with (root_path, value, parent)
+        and return the style.
+      key_color: The color for label-style keys under this extension. It can be
+        a tuple of (color, background-color) or a function that takes
+        (root_path, value, parent) and returns the color tuple.
+      include_keys: A list of keys to include when displaying the sub-nodes of
+        the object. If None, all keys will be displayed. If a function, it will
+        be called with (root_path, value, parent) and return whether the key
+        should be included.
+      exclude_keys: A set of keys to exclude when displaying the sub-nodes of
+        the object. If None, all keys will be displayed. If a function, it will
+        be called with (root_path, value, parent) and return whether the key
+        should be excluded.
+      enable_key_tooltip: Whether to enable the tooltip for the object name.
+      collapse_level: The level of collapsing. If 0, the object will be
+        collapsed (without showing its sub-nodes). If 1, the immediate sub-nodes
+        will be shown in collapsed form. If None, all sub-tree will be shown.
+      uncollapse: Indivdual nodes to uncollapse. It can be a KeyPathSet or a
+        function that takes (root_path, value, parent) and returns a KeyPathSet.
+      child_config: The configs for the immediate child nodes of the object
+        being rendered. It's a dictionary of (key, child-config) where the key
+        is the name of the child node and the child-config is a dictionary of
+        (key, value) to override the default configs for the child node.
+      extra_flags: A dictionary of user-defined flags to control the rendering
+        behavior.
+      highlight: A function that takes (root_path, value, parent) and returns
+        whether the node should be highlighted.
+      lowlight: A function that takes (root_path, value, parent) and returns
+        whether the node should be lowlighted.
+      debug: Whether to show debug information for this rendering.
 
     Returns:
       The rendered HTML.
     """
     root_path = root_path or KeyPath()
+    child_config = child_config or {}
+    extra_flags = extra_flags or {}
     uncollapse = self.init_uncollapse(uncollapse)
 
-    child_collapse_level = collapse_level
-    if isinstance(value, HtmlTreeView.Extension):
-      subtree_uncollapse_level = value._html_tree_view_uncollapse_level()    # pylint: disable=protected-access
-
-      # If the extension has child levels to uncollapse, honor them above the
-      # collapse level passed from the root. However, we can see the
-      # uncollapsed extension subtree only when the extension's parent node is
-      # uncollapsed.
-      child_collapse_level = self.max_collapse_level(
-          collapse_level, subtree_uncollapse_level, root_path
-      )
-      if not callable(uncollapse):
-        extension_uncollapse = value._html_tree_view_uncollapse().copy()  # pylint: disable=protected-access
-        if extension_uncollapse:
-          extension_uncollapse.rebase(root_path)
-          uncollapse = uncollapse.union(extension_uncollapse)
-
     summary = self.summary(
         value,
         name=name,
         parent=parent,
         root_path=root_path,
+        css_classes=css_classes,
         title=title,
+        summary_color=summary_color,
         enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
         enable_summary_tooltip=enable_summary_tooltip,
+        enable_key_tooltip=enable_key_tooltip,
         max_summary_len_for_str=max_summary_len_for_str,
-        **kwargs,
+        extra_flags=extra_flags,
     )
+
+    if debug:
+      debug_info = Html.element(
+          'div',
+          [
+              Html.element(
+                  'span', ['DEBUG'], css_classes=['debug-info-trigger']
+              ),
+              self.tooltip(
+                  dict(
+                      # Most error-prone settings.
+                      css_classes=css_classes,
+                      collapse_level=collapse_level,
+                      uncollapse=uncollapse,
+                      extra_flags=extra_flags,
+                      child_config=child_config,
+                      # Relative obvious settings.
+                      key_style=key_style,
+                      key_color=key_color,
+                      include_keys=include_keys,
+                      exclude_keys=exclude_keys,
+                      # More obvious settings.
+                      summary_color=summary_color,
+                      enable_summary=enable_summary,
+                      enable_summary_for_str=enable_summary_for_str,
+                      max_summary_len_for_str=max_summary_len_for_str,
+                      enable_summary_tooltip=enable_summary_tooltip,
+                      enable_key_tooltip=enable_key_tooltip,
+                  ),
+                  name='debug_info',
+                  parent=parent,
+                  root_path=root_path,
+                  css_classes=['debug-info'],
+              ),
+          ],
+      ).add_style(
+          """
+          .debug-info-trigger {
+            display: inline-flex;
+            cursor: pointer;
+            font-size: 0.6em;
+            background-color: red;
+            color: white;
+            padding: 5px;
+            border-radius: 3px;
+            margin: 5px 0 5px 0;
+          }
+          .debug-info-trigger:hover + span.tooltip {
+            visibility: visible;
+          }
+          """
+      )
+    else:
+      debug_info = None
+
     content = self.content(
         value,
         name=name,
         parent=parent,
         root_path=root_path,
-        filter=filter,
-        special_keys=special_keys,
-        include_keys=include_keys,
-        exclude_keys=exclude_keys,
-        collapse_level=child_collapse_level,
-        uncollapse=uncollapse,
+        css_classes=css_classes if summary is None else None,
+        # Summary settings (child nodes).
+        enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
         max_summary_len_for_str=max_summary_len_for_str,
         enable_summary_tooltip=enable_summary_tooltip,
+        # Content settings.
+        key_style=key_style,
+        key_color=key_color,
         enable_key_tooltip=enable_key_tooltip,
-        **kwargs,
-    )
-    extension_style = (
-        value._html_style() if isinstance(value, HtmlView.Extension) else []  # pylint: disable=protected-access
+        include_keys=include_keys,
+        exclude_keys=exclude_keys,
+        collapse_level=collapse_level,
+        uncollapse=uncollapse,
+        highlight=highlight,
+        lowlight=lowlight,
+        child_config=child_config,
+        extra_flags=extra_flags,
+        debug=debug,
     )
 
     if summary is None:
       content = Html.from_value(content)
       assert content is not None
-      content.add_style(*extension_style)
-      return content
+      return debug_info + content
 
-    collapse_view = self.should_collapse(
-        value, parent=parent, root_path=root_path,
+    collapse_details = self.should_collapse(
+        value, name=name, parent=parent, root_path=root_path,
         collapse_level=collapse_level, uncollapse=uncollapse,
     )
     return Html.element(
         'details',
         [
             summary,
+            debug_info,
             content,
         ],
-        options=[None if collapse_view else 'open'],
-        css_class=[
+        options=[None if collapse_details else 'open'],
+        css_classes=[
             'pyglove',
             self.css_class_name(value),
+            css_classes,
         ],
     ).add_style(
         """
@@ -435,10 +427,7 @@ class HtmlTreeView(HtmlView):
           border: 1px solid #aaa;
           border-radius: 4px;
           padding: 0.5em 0.5em 0;
-          margin: 0.1em 0;
-        }
-        details.pyglove.special_value {
-          margin-bottom: 0.75em;
+          margin: 0.25em 0;
         }
         details.pyglove[open] {
           padding: 0.5em 0.5em 0.5em;
@@ -450,38 +439,44 @@ class HtmlTreeView(HtmlView):
           opacity: 0.2;
         }
         """,
-        *extension_style,
     )
 
-  def init_uncollapse(
-      self,
-      uncollapse: Union[
-          Iterable[Union[KeyPath, str]], base.NodeFilter, None
-      ] = HtmlView.PresetArgValue(None),
-  ) -> Union[KeyPathSet, base.NodeFilter]:
-    """Normalize the uncollapse argument."""
-    if uncollapse is None:
-      return KeyPathSet()
-    elif callable(uncollapse):
-      return uncollapse
-    else:
-      return KeyPathSet.from_value(uncollapse, include_intermediate=True)
-
   def should_collapse(
       self,
       value: Any,
+      name: Optional[str],
       root_path: KeyPath,
       parent: Any,
-      collapse_level: Optional[int] = 0,
+      collapse_level: Optional[int] = 1,
       uncollapse: Union[KeyPathSet, base.NodeFilter] = None,
   ) -> bool:
-    """Returns whether the object should be collapsed."""
-    if collapse_level is None or root_path.depth < collapse_level:
+    """Returns True if the object should be collapsed.
+
+    Args:
+      value: The value to render.
+      name: The referred field name of the value.
+      root_path: The root path of the value.
+      parent: The parent of the value.
+      collapse_level: The level of collapsing. If 0, the object will be
+        collapsed (without showing its sub-nodes). If 1, the immediate sub-nodes
+        will be shown in collapsed form. If None, all sub-tree will be shown.
+      uncollapse: Indivdual nodes to uncollapse. It can be a KeyPathSet or a
+        function that takes (root_path, value, parent) and returns a KeyPathSet.
+
+    Returns:
+      True if the object should be collapsed.
+    """
+    if collapse_level is None or collapse_level > 0:
       return False
     if callable(uncollapse):
       return not uncollapse(root_path, value, parent)
-    else:
-      return root_path not in uncollapse
+    if root_path in uncollapse:
+      return False
+    # Always uncollapse simple types.
+    if (name is not None
+        and isinstance(value, (bool, int, float, str, type(None)))):
+      return False
+    return True
 
   def needs_summary(
       self,
@@ -490,21 +485,38 @@ class HtmlTreeView(HtmlView):
       name: Optional[str] = None,
       parent: Any = None,
       title: Union[str, Html, None] = None,
-      enable_summary: Optional[bool] = HtmlView.PresetArgValue(True),
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
+      enable_summary: Optional[bool] = None,
+      enable_summary_for_str: bool = True,
+      max_summary_len_for_str: int = 80,
   ) -> bool:
-    """Returns whether the object needs a summary."""
+    """Returns True if the object needs a summary.
+
+    Args:
+      value: The value to render.
+      name: The referred field name of the value.
+      parent: The parent of the value.
+      title: The title of the summary.
+      enable_summary: Whether to enable the summary. If None, summary will
+        be enabled for complex types or when string exceeds
+        `max_summary_len_for_str`.
+      enable_summary_for_str: Whether to enable the summary for strings.
+      max_summary_len_for_str: The maximum length of the string to display.
+
+    Returns:
+      True if the object needs a summary.
+    """
     del parent
-    if enable_summary is None:
-      if name is not None or title is not None or not (
-          isinstance(value, (int, float, bool, type(None)))
-          or (
-              isinstance(value, str)
-              and len(value) <= max_summary_len_for_str
-          )
-      ):
-        return True
-    return enable_summary
+    if isinstance(enable_summary, bool):
+      return enable_summary
+    assert enable_summary is None
+    if not enable_summary_for_str and isinstance(value, str):
+      return False
+    if name is None and title is None and (
+        isinstance(value, (int, float, bool, type(None)))
+        or (isinstance(value, str) and len(value) <= max_summary_len_for_str)
+    ):
+      return False
+    return True
 
   @HtmlView.extension_method('_html_tree_view_summary')
   def summary(
@@ -514,31 +526,52 @@ class HtmlTreeView(HtmlView):
       name: Optional[str] = None,
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
+      css_classes: Optional[Sequence[str]] = None,
       title: Union[str, Html, None] = None,
-      enable_summary: Optional[bool] = HtmlView.PresetArgValue(None),
-      enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-      **kwargs
+      enable_summary: Optional[bool] = None,
+      enable_summary_tooltip: bool = True,
+      summary_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      max_summary_len_for_str: int = 80,
+      enable_summary_for_str: bool = True,
+      enable_key_tooltip: bool = True,
+      summary_tooltip_fn: Optional[Callable[..., Html]] = None,
+      key_tooltip_fn: Optional[Callable[..., Html]] = None,
+      extra_flags: Optional[Dict[str, Any]] = None,
   ) -> Optional[Html]:
-    """Renders a summary for the value.
+    """Renders the summary for an input value.
 
     Args:
       value: The value to render.
-      name: The name of the value.
+      name: The referred field name of the value.
       parent: The parent of the value.
       root_path: The root path of the value.
+      css_classes: The CSS classes to add to the HTML element.
       title: The title of the summary.
       enable_summary: Whether to enable the summary. If None, summary will
         be enabled for complex types or when string exceeds
         `max_summary_len_for_str`.
-      enable_summary_tooltip: Whether to enable the tooltip for the summary.
+      enable_summary_tooltip: Whether to enable the summary tooltip.
+      summary_color: The color of the summary. If None, the summary will be
+        rendered without a color. If a tuple, the first element is the text
+        color and the second element is the background color. If a function,
+        the function takes (root_path, value, parent) and returns a tuple of
+        (text_color, background_color).
       max_summary_len_for_str: The maximum length of the string to display.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+      enable_summary_for_str: Whether to enable the summary for strings.
+      enable_key_tooltip: Whether to enable the key tooltip.
+      summary_tooltip_fn: The function to render the summary tooltip.
+      key_tooltip_fn: The function to render the key tooltip.
+      extra_flags: The extra flags to pass to the summary.
 
     Returns:
       An optional HTML object representing the summary of the value. If None,
-      the summary will be hidden.
+      the summary will not be rendered.
     """
+    del extra_flags
+    root_path = root_path or KeyPath()
     if not self.needs_summary(
         value,
         name=name,
@@ -546,16 +579,27 @@ class HtmlTreeView(HtmlView):
         title=title,
         max_summary_len_for_str=max_summary_len_for_str,
         enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
     ):
       return None
 
+    key_tooltip_fn = key_tooltip_fn or self.tooltip
+    summary_tooltip_fn = summary_tooltip_fn or self.tooltip
+
     def make_title(value: Any):
-      if isinstance(value, str):
-        if len(value) > max_summary_len_for_str:
-          value = value[:max_summary_len_for_str] + '...'
-        return Html.escape(repr(value))
+      if inspect.isclass(value):
+        return 'type'
+      elif isinstance(value, (int, float, bool, str)):
+        return type(value).__name__
       return f'{type(value).__name__}(...)'
 
+    if name is not None:
+      summary_color = self.get_color(
+          summary_color, root_path + name, value, parent
+      )
+    else:
+      summary_color = (None, None)
+
     return Html.element(
         'summary',
         [
@@ -564,8 +608,19 @@ class HtmlTreeView(HtmlView):
                 'div',
                 [
                     name,
+                    key_tooltip_fn(   # pylint: disable=g-long-ternary
+                        root_path,
+                        name=name,
+                        parent=parent,
+                        root_path=root_path,
+                        css_classes=css_classes,
+                    ) if enable_key_tooltip else None,
                 ],
-                css_class=['summary_name']
+                css_classes=['summary-name', css_classes],
+                styles=dict(
+                    color=summary_color[0],
+                    background_color=summary_color[1],
+                ),
             ) if name else None,
 
             # Summary title
@@ -574,16 +629,15 @@ class HtmlTreeView(HtmlView):
                 [
                     title or make_title(value),
                 ],
-                css_class=['summary_title']
+                css_classes=['summary-title', css_classes],
             ),
 
-            # Tooltip.
-            lambda: self.tooltip(   # pylint: disable=g-long-ternary
+            # Summary tooltip.
+            lambda: summary_tooltip_fn(   # pylint: disable=g-long-ternary
                 value,
-                name=name,
                 parent=parent,
                 root_path=root_path,
-                **kwargs,
+                css_classes=css_classes,
             ) if enable_summary_tooltip else None,
         ],
     ).add_style(
@@ -594,24 +648,25 @@ class HtmlTreeView(HtmlView):
           margin: -0.5em -0.5em 0;
           padding: 0.5em;
         }
-        .summary_name {
+        .summary-name {
           display: inline;
-          padding: 0 5px;
+          padding: 3px 5px 3px 5px;
+          margin: 0 5px;
+          border-radius: 3px;
         }
-        .summary_title {
+        .summary-title {
           display: inline;
         }
-        .summary_name + div.summary_title {
+        .summary-name + div.summary-title {
           display: inline;
           color: #aaa;
         }
-        .summary_title:hover + span.tooltip {
+        .summary-title:hover + span.tooltip {
           visibility: visible;
         }
-        /* Type-specific styles. */
-        .pyglove.str .summary_title {
-          color: darkred;
-          font-style: italic;
+        .summary-name:hover > span.tooltip {
+          visibility: visible;
+          background-color: darkblue;
         }
         """
     )
@@ -621,80 +676,90 @@ class HtmlTreeView(HtmlView):
   # delegated to `HtmlTreeView.Extension`.
   def object_key(
       self,
-      key: Union[str, int],
+      root_path: KeyPath,
       *,
-      name: Optional[str] = None,
+      value: Any,
       parent: Any,
-      root_path: Optional[KeyPath] = None,
-      css_class: Union[str, Sequence[str], None] = None,
-      key_color: Optional[str] = None,
-      enable_tooltip: bool = HtmlView.PresetArgValue(True),
-      **kwargs
+      css_classes: Optional[Sequence[str]] = None,
+      key_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      enable_key_tooltip: bool = True,
+      key_tooltip_fn: Optional[Callable[..., Html]] = None,
+      **kwargs,
   ) -> Html:
-    """Renders the key for the value.
+    """Renders a label-style key for the value.
 
     Args:
-      key: The key of the value.
-      name: The name of the value.
-      parent: The parent value of the key.
       root_path: The root path of the value.
-      css_class: Additional CSS classes to add to the HTML element.
-      key_color: The color of the key.
-      enable_tooltip: Whether to enable the tooltip.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+      value: The value to render.
+      parent: The parent of the value.
+      css_classes: The CSS classes to add to the HTML element.
+      key_color: The color of the key. If None, the key will be rendered
+        without a color. If a tuple, the first element is the text color and
+        the second element is the background color. If a function, the function
+        takes (root_path, value, parent) and returns a tuple of (text_color,
+        background_color).
+      enable_key_tooltip: Whether to enable the tooltip.
+      key_tooltip_fn: The function to render the key tooltip.
+      **kwargs: Additional arguments passed by the user that will be ignored.
 
     Returns:
       The rendered HTML as the key of the value.
     """
+    del kwargs
+    key_tooltip_fn = key_tooltip_fn or self.tooltip
+    key_color = self.get_color(key_color, root_path, value, parent)
     return (
         # Key span.
         Html.element(
             'span',
             [
-                str(key),
+                str(root_path.key),
             ],
-            css_class=[
-                'object_key',
-                type(key).__name__,
-                css_class,
+            css_classes=[
+                'object-key',
+                type(root_path.key).__name__,
+                css_classes,
             ],
-            style=dict(
-                color=key_color,
+            styles=dict(
+                color=key_color[0],
+                background_color=key_color[1],
             )
         ) + (
             # Tooltip if enabled.
-            lambda: self.tooltip(  # pylint: disable=g-long-ternary
+            lambda: key_tooltip_fn(  # pylint: disable=g-long-ternary
                 value=root_path,
                 root_path=root_path,
-                name=name,
                 parent=parent,
-                **kwargs
-            ) if enable_tooltip else None
+            ) if enable_key_tooltip else None
         )
     ).add_style(
         """
         /* Object key styles. */
-        .object_key {
-          margin-right: 0.25em;
+        .object-key {
+          margin: 0.15em 0.3em 0.15em 0;
+          display: block;
         }
-        .object_key:hover + .tooltip {
+        .object-key:hover + .tooltip {
           visibility: visible;
           background-color: darkblue;
         }
-        .object_key.str {
+        .object-key.str {
           color: gray;
           border: 1px solid lightgray;
           background-color: ButtonFace;
           border-radius: 0.2em;
           padding: 0.3em;
         }
-        .object_key.int::before{
+        .object-key.int::before{
           content: '[';
         }
-        .object_key.int::after{
+        .object-key.int::after{
           content: ']';
         }
-        .object_key.int{
+        .object-key.int{
           border: 0;
           color: lightgray;
           background-color: transparent;
@@ -712,20 +777,41 @@ class HtmlTreeView(HtmlView):
       name: Optional[str] = None,
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
-      special_keys: Optional[Sequence[Union[int, str]]] = None,
-      include_keys: Optional[Iterable[Union[int, str]]] = None,
-      exclude_keys: Optional[Iterable[Union[int, str]]] = None,
-      collapse_level: Optional[int] = HtmlView.PresetArgValue(1),
-      uncollapse: Union[
-          KeyPathSet, base.NodeFilter, None
-      ] = HtmlView.PresetArgValue(None),
-      filter: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),  # pylint: disable=redefined-builtin
-      highlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      lowlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-      enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-      enable_key_tooltip: bool = HtmlView.PresetArgValue(True),
-      **kwargs
+      css_classes: Optional[Sequence[str]] = None,
+      # Summary settings (for child nodes).
+      enable_summary: Optional[bool] = None,
+      enable_summary_for_str: bool = True,
+      max_summary_len_for_str: int = 80,
+      enable_summary_tooltip: bool = True,
+      # Content settings.
+      key_style: Union[
+          Literal['label', 'summary'],
+          Callable[[KeyPath, Any, Any], Literal['label', 'summary']]
+      ] = 'summary',
+      key_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      include_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      exclude_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      enable_key_tooltip: bool = True,
+      # Collapse settings.
+      collapse_level: Optional[int] = 1,
+      uncollapse: Union[KeyPathSet, base.NodeFilter, None] = None,
+      # Other settings.
+      highlight: Optional[base.NodeFilter] = None,
+      lowlight: Optional[base.NodeFilter] = None,
+      child_config: Optional[Dict[str, Any]] = None,
+      extra_flags: Optional[Dict[str, Any]] = None,
+      debug: bool = False,
   ) -> Html:
     """Renders the main content for the value.
 
@@ -734,27 +820,42 @@ class HtmlTreeView(HtmlView):
       name: The name of the value.
       parent: The parent of the value.
       root_path: The root path of the value.
-      special_keys: The special keys to display (at the immediate child level).
-      include_keys: The keys to include (at the immediate child level).
-      exclude_keys: The keys to exclude (at the immediate child level).
-      collapse_level: The level to collapse the tree (relative to this node).
+      css_classes: CSS classes to add to the HTML element.
+      enable_summary: Whether to enable the summary.
+      enable_summary_for_str: Whether to enable the summary for string.
+      max_summary_len_for_str: The maximum length of the string to display.
+      enable_summary_tooltip: Whether to enable the summary tooltip.
+      key_style: The style of the key. It can be either 'label' or 'summary'.
+        If it is a function, the function takes (root_path, value, parent) and
+        returns either 'label' or 'summary'.
+      key_color: The color of the key. If it is a tuple, the first element is
+        the text color and the second element is the background color. If it is
+        a function, the function takes (root_path, value, parent) and returns
+        a tuple of (text_color, background_color).
+      include_keys: The keys to include (at the immediate child level). If it is
+        a function, the function takes (root_path, value, parent) and returns
+        an iterable of keys to include.
+      exclude_keys: The keys to exclude (at the immediate child level). If it is
+        a function, the function takes (root_path, value, parent) and returns
+        an iterable of keys to exclude.
+      enable_key_tooltip: Whether to enable the key tooltip.
+      collapse_level: The level to collapse the tree.
       uncollapse: A key path set (relative to root_path) for the nodes to
         uncollapse. or a function with signature (path, value, parent) -> bool
         to filter nodes to uncollapse.
-      filter: A function with signature (path, value, parent) -> include
-        to determine whether to include a field (at all levels).
       highlight: A function with signature (path, value, parent) -> bool
-        to determine whether to highlight a field (at all levels).
+        to determine whether to highlight.
       lowlight: A function with signature (path, value, parent) -> bool
-        to determine whether to lowlight a field (at all levels).
-      max_summary_len_for_str: The maximum length of the string to display.
-      enable_summary_tooltip: Whether to enable the summary tooltip.
-      enable_key_tooltip: Whether to enable the key tooltip.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+        to determine whether to lowlight.
+      child_config: The configuration for rendering the child nodes.
+      extra_flags: Extra flags to pass to the child render.
+      debug: Whether to enable debug mode.
 
     Returns:
       The rendered HTML as the main content of the value.
     """
+    root_path = root_path or KeyPath()
+
     if isinstance(value, (tuple, list)):
       items = {i: v for i, v in enumerate(value)}
     elif isinstance(value, dict):
@@ -762,25 +863,31 @@ class HtmlTreeView(HtmlView):
     else:
       return self.simple_value(
           value, name=name, parent=parent, root_path=root_path,
+          css_classes=css_classes,
           max_summary_len_for_str=max_summary_len_for_str
       )
     return self.complex_value(
         items,
         name=name,
         parent=value,
-        root_path=root_path or KeyPath(),
-        special_keys=special_keys,
+        root_path=root_path,
+        css_classes=css_classes,
+        enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
+        max_summary_len_for_str=max_summary_len_for_str,
+        enable_summary_tooltip=enable_summary_tooltip,
+        key_style=key_style,
+        key_color=key_color,
+        enable_key_tooltip=enable_key_tooltip,
         include_keys=include_keys,
         exclude_keys=exclude_keys,
         collapse_level=collapse_level,
         uncollapse=uncollapse,
-        filter=filter,
+        child_config=child_config,
         highlight=highlight,
         lowlight=lowlight,
-        max_summary_len_for_str=max_summary_len_for_str,
-        enable_summary_tooltip=enable_summary_tooltip,
-        enable_key_tooltip=enable_key_tooltip,
-        **kwargs,
+        extra_flags=extra_flags,
+        debug=debug,
     )
 
   def simple_value(
@@ -790,8 +897,8 @@ class HtmlTreeView(HtmlView):
       name: Optional[str] = None,
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
-      css_class: Union[str, Sequence[str], None] = None,
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
+      css_classes: Optional[Sequence[str]] = None,
+      max_summary_len_for_str: int = 80,
   ) -> Html:
     """Renders a simple value.
 
@@ -800,7 +907,7 @@ class HtmlTreeView(HtmlView):
       name: The name of the value.
       parent: The parent of the value.
       root_path: The root path of the value.
-      css_class: Additional CSS classes to add to the HTML element.
+      css_classes: CSS classes to add to the HTML element.
       max_summary_len_for_str: The maximum length of the string to display.
 
     Returns:
@@ -824,26 +931,26 @@ class HtmlTreeView(HtmlView):
         [
             Html.escape(value_repr),
         ],
-        css_class=[
-            'simple_value',
+        css_classes=[
+            'simple-value',
             self.css_class_name(value),
-            css_class,
+            css_classes,
         ],
     ).add_style(
         """
         /* Simple value styles. */
-        .simple_value {
+        .simple-value {
           color: blue;
           display: inline-block;
           white-space: pre-wrap;
           padding: 0.2em;
           margin-top: 0.15em;
         }
-        .simple_value.str {
+        .simple-value.str {
           color: darkred;
           font-style: italic;
         }
-        .simple_value.int, .simple_value.float {
+        .simple-value.int, .simple-value.float {
           color: darkblue;
         }
         """
@@ -856,55 +963,86 @@ class HtmlTreeView(HtmlView):
       parent: Any,
       root_path: KeyPath,
       name: Optional[str] = None,
-      css_class: Union[str, Sequence[str], None] = None,
+      css_classes: Optional[Sequence[str]] = None,
+      # Summary settings (for child nodes).
+      enable_summary: Optional[bool] = None,
+      enable_summary_for_str: bool = True,
+      max_summary_len_for_str: int = 80,
+      enable_summary_tooltip: bool = True,
+      # Content settings.
+      key_style: Union[
+          Literal['label', 'summary'],
+          Callable[..., Literal['label', 'summary']]
+      ] = 'summary',
+      key_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = None,
+      include_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      exclude_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = None,
+      enable_key_tooltip: bool = True,
+      # Collapse settings.
+      collapse_level: Optional[int] = 1,
+      uncollapse: Union[KeyPathSet, base.NodeFilter, None] = None,
+      # Other settings.
+      child_config: Optional[Dict[str, Any]] = None,
+      highlight: Optional[base.NodeFilter] = None,
+      lowlight: Optional[base.NodeFilter] = None,
+      # Custom render functions.
       render_key_fn: Optional[Callable[..., Html]] = None,
       render_value_fn: Optional[Callable[..., Html]] = None,
-      special_keys: Optional[Sequence[Union[int, str]]] = None,
-      include_keys: Optional[Iterable[Union[int, str]]] = None,
-      exclude_keys: Optional[Iterable[Union[int, str]]] = None,
-      collapse_level: Optional[int] = HtmlView.PresetArgValue(1),
-      uncollapse: Union[
-          KeyPathSet, base.NodeFilter, None
-      ] = HtmlView.PresetArgValue(None),
-      filter: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),  # pylint: disable=redefined-builtin
-      highlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      lowlight: Optional[base.NodeFilter] = HtmlView.PresetArgValue(None),
-      max_summary_len_for_str: int = HtmlView.PresetArgValue(40),
-      enable_summary_tooltip: bool = HtmlView.PresetArgValue(True),
-      enable_key_tooltip: bool = HtmlView.PresetArgValue(True),
-      **kwargs,
+      extra_flags: Optional[Dict[str, Any]] = None,
+      debug: bool = False,
   ) -> Html:
     """Renders a list of key-value pairs.
 
     Args:
       kv: The key-value pairs to render.
-      parent: The parent value of the key-value pairs.
+      parent: The parent of the value.
       root_path: The root path of the value.
       name: The name of the value.
-      css_class: Additional CSS classes to add to the HTML element.
-      render_key_fn: A function to render the key. The function has the
-        same signature as `HtmlTreeView.object_key`.
-        If None, `HtmlTreeView.object_key` will be used to render the key.
-      render_value_fn: A function to render the value. The function has the
-        same signature as `HtmlTreeView.render`.
-        If None, `HtmlTreeView.render` will be used to render child value.
-      special_keys: The special keys to display (at the immediate child level).
-      include_keys: The keys to include (at the immediate child level).
-      exclude_keys: The keys to exclude (at the immediate child level).
-      collapse_level: The level to collapse the tree (relative to this node).
+      css_classes: CSS classes to add to the HTML element.
+      enable_summary: Whether to enable the summary. If None, the default is
+        to enable the summary for non-string and disable the summary for
+        string.
+      enable_summary_for_str: Whether to enable the summary for string.
+      max_summary_len_for_str: The maximum length of the string to display.
+      enable_summary_tooltip: Whether to enable the summary tooltip.
+      key_style: The style of the key. It can be either 'label' or 'summary'.
+        If it is a function, the function takes (root_path, value, parent) and
+        returns either 'label' or 'summary'.
+      key_color: The color of the key. If it is a tuple, the first element is
+        the text color and the second element is the background color. If it is
+        a function, the function takes (root_path, value, parent) and returns
+        a tuple of (text_color, background_color).
+      include_keys: The keys to include (at the immediate child level). If it is
+        a function, the function takes (root_path, value, parent) and returns
+        an iterable of keys to include.
+      exclude_keys: The keys to exclude (at the immediate child level). If it is
+        a function, the function takes (root_path, value, parent) and returns
+        an iterable of keys to exclude.
+      enable_key_tooltip: Whether to enable the key tooltip.
+      collapse_level: The level to collapse the tree.
       uncollapse: A key path set (relative to root_path) for the nodes to
         uncollapse. or a function with signature (path, value, parent) -> bool
         to filter nodes to uncollapse.
-      filter: A function with signature (path, value, parent) -> include
-        to determine whether to include.
+      child_config: The configuration for rendering the child nodes.
       highlight: A function with signature (path, value, parent) -> bool
         to determine whether to highlight.
       lowlight: A function with signature (path, value, parent) -> bool
         to determine whether to lowlight.
-      max_summary_len_for_str: The maximum length of the string to display.
-      enable_summary_tooltip: Whether to enable the summary tooltip.
-      enable_key_tooltip: Whether to enable the key tooltip.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+      render_key_fn: A custom function to render the label-style key.
+      render_value_fn: A custom function to render the child value.
+      extra_flags: Extra flags to pass to the child render.
+      debug: Whether to enable debug mode.
 
     Returns:
       The rendered HTML as the key-value pairs.
@@ -912,142 +1050,156 @@ class HtmlTreeView(HtmlView):
     del name
     root_path = root_path or KeyPath()
     uncollapse = self.init_uncollapse(uncollapse)
+    extra_flags = extra_flags or {}
 
-    if isinstance(parent, HtmlTreeView.Extension):
-      special_keys = special_keys or parent._html_tree_view_special_keys()  # pylint: disable=protected-access
-      include_keys = include_keys or parent._html_tree_view_include_keys()  # pylint: disable=protected-access
-      exclude_keys = exclude_keys or parent._html_tree_view_exclude_keys()  # pylint: disable=protected-access
+    inherited_kwargs = dict(
+        # For child summary.
+        enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
+        max_summary_len_for_str=max_summary_len_for_str,
+        enable_summary_tooltip=enable_summary_tooltip,
+        # For child content.
+        enable_key_tooltip=enable_key_tooltip,
+        key_style=key_style,
+        key_color=key_color,
+        include_keys=include_keys if callable(include_keys) else None,
+        exclude_keys=exclude_keys if callable(exclude_keys) else None,
+        collapse_level=None if collapse_level is None else (collapse_level - 1),
+        uncollapse=uncollapse,
+        highlight=highlight,
+        lowlight=lowlight,
+        extra_flags=extra_flags,
+        debug=debug,
+    )
 
-    special_keys = special_keys or []
-    include_keys = set(include_keys or [])
-    exclude_keys = set(exclude_keys or [])
+    render_key_fn = render_key_fn or HtmlTreeView.object_key
+    render_value_fn = render_value_fn or HtmlTreeView.render
 
-    render_key_fn = render_key_fn or self.object_key
-    render_value_fn = render_value_fn or self.render
+    def render_child_key(child_path, value, parent, child_kwargs):
+      render_child_key_fn = child_kwargs['extra_flags'].get(
+          'render_key_fn', render_key_fn
+      )
+      return render_child_key_fn(
+          self,
+          child_path,
+          value=value,
+          parent=parent,
+          **child_kwargs
+      )
 
+    def render_child_value(name, value, child_path, child_kwargs):
+      render_child_value_fn = child_kwargs['extra_flags'].get(
+          'render_value_fn', render_value_fn
+      )
+      child_html = render_child_value_fn(
+          self,
+          value=value, name=child_kwargs.pop('name', name),
+          parent=parent,
+          root_path=child_path,
+          **child_kwargs
+      )
+      should_highlight = highlight and highlight(child_path, value, parent)
+      should_lowlight = lowlight and lowlight(child_path, value, parent)
+      if should_highlight or should_lowlight:
+        return Html.element(
+            'div', [child_html],
+            css_classes=[
+                'highlight' if should_highlight else None,
+                'lowlight' if should_lowlight else None,
+            ],
+        )
+      else:
+        return child_html
+
+    has_child = False
     s = Html()
     if kv:
-      include_keys = include_keys or set(kv.keys())
-      if filter is not None:
-        include_keys -= set(
-            k for k, v in kv.items()
-            if not filter(root_path + k, v, parent)
-        )
-      if exclude_keys:
-        include_keys -= exclude_keys
-
-      if special_keys:
-        for k in special_keys:
-          if k in include_keys and k in kv:
-            child_path = root_path + k
-            v = kv[k]
+      # Compute included keys.
+      if callable(include_keys):
+        include_keys = [
+            k for k, v in kv.items() if include_keys(root_path + k, v, parent)
+        ]
+      elif include_keys is not None:
+        include_keys = list(k for k in include_keys if k in kv)
+      else:
+        include_keys = list(kv.keys())
+
+      # Filter with excluded keys.
+      if callable(exclude_keys):
+        include_keys = [
+            k for k in include_keys if not exclude_keys(
+                root_path + k, kv[k], parent
+            )
+        ]
+      elif exclude_keys is not None:
+        exclude_keys = set(exclude_keys)
+        include_keys = [k for k in include_keys if k not in exclude_keys]
+
+      # Figure out keys of different styles.
+      label_keys = []
+      summary_keys = []
+      if isinstance(parent, (tuple, list)) or key_style == 'label':
+        label_keys = include_keys
+      elif key_style == 'summary':
+        summary_keys = include_keys
+      else:
+        assert callable(key_style), key_style
+        for k in include_keys:
+          ks = key_style(root_path + k, kv[k], parent)
+          if ks == 'summary':
+            summary_keys.append(k)
+          elif ks == 'label':
+            label_keys.append(k)
+
+      # Render child nodes with summary keys.
+      if summary_keys:
+        for k in summary_keys:
+          child_path = root_path + k
+          child_kwargs = self.get_child_kwargs(
+              inherited_kwargs, child_config, k, root_path
+          )
+          s.write(render_child_value(k, kv[k], child_path, child_kwargs))
+          has_child = True
+
+      # Render child nodes with label keys.
+      if label_keys:
+        s.write('<table>')
+        for k in label_keys:
+          v = kv[k]
+          child_path = root_path + k
+          child_kwargs = self.get_child_kwargs(
+              inherited_kwargs, child_config, k, root_path
+          )
+          key_cell = render_child_key(child_path, v, parent, child_kwargs)
+          value_cell = render_child_value(None, v, child_path, child_kwargs)
+          if value_cell is not None:
             s.write(
                 Html.element(
-                    'div',
+                    'tr',
                     [
-                        render_value_fn(
-                            value=v,
-                            name=k,
-                            parent=parent,
-                            root_path=child_path,
-                            filter=filter,
-                            special_keys=None,
-                            include_keys=None,
-                            exclude_keys=None,
-                            collapse_level=collapse_level,
-                            uncollapse=uncollapse,
-                            highlight=highlight,
-                            lowlight=lowlight,
-                            max_summary_len_for_str=max_summary_len_for_str,
-                            enable_summary_tooltip=enable_summary_tooltip,
-                            enable_key_tooltip=enable_key_tooltip,
-                            **kwargs
-                        )
-                    ],
-                    css_class=[
-                        'special_value',
-                        (
-                            'highlight' if highlight
-                            and highlight(child_path, v, parent) else None
-                        ),
-                        (
-                            'lowlight' if lowlight
-                            and lowlight(child_path, v, parent) else None
-                        )
+                        '<td>', key_cell, '</td>',
+                        '<td>', value_cell, '</td>',
                     ],
                 )
             )
-            include_keys.remove(k)
-
-      if include_keys:
-        s.write('<table>')
-        for k, v in kv.items():
-          if k not in include_keys:
-            continue
-          child_path = root_path + k
-          key_cell = render_key_fn(
-              key=k,
-              parent=parent,
-              root_path=child_path,
-              enable_tooltip=enable_key_tooltip,
-          )
-          value_cell = Html.element(
-              'div',
-              [
-                  render_value_fn(
-                      value=v,
-                      name=None,
-                      parent=parent,
-                      root_path=child_path,
-                      special_keys=None,
-                      include_keys=None,
-                      exclude_keys=None,
-                      collapse_level=collapse_level,
-                      uncollapse=uncollapse,
-                      filter=filter,
-                      highlight=highlight,
-                      lowlight=lowlight,
-                      max_summary_len_for_str=max_summary_len_for_str,
-                      enable_summary_tooltip=enable_summary_tooltip,
-                      enable_key_tooltip=enable_key_tooltip,
-                      **kwargs,
-                  )
-              ],
-              css_class=[
-                  (
-                      'highlight' if highlight
-                      and highlight(child_path, v, parent) else None
-                  ),
-                  (
-                      'lowlight' if lowlight
-                      and lowlight(child_path, v, parent) else None
-                  )
-              ],
-          )
-          s.write(
-              Html.element(
-                  'tr',
-                  [
-                      '<td>', key_cell, '</td>',
-                      '<td>', value_cell, '</td>',
-                  ],
-              )
-          )
+            has_child = True
         s.write('</table>')
-    else:
-      s.write(Html.element('span', css_class=['empty_container']))
+
+    if not has_child:
+      s.write(Html.element('span', css_classes=['empty-container']))
+
     return Html.element(
         'div',
         [s],
-        css_class=[
-            'complex_value',
+        css_classes=[
+            'complex-value',
             self.css_class_name(parent),
-            css_class,
+            css_classes,
         ]
     ).add_style(
         """
         /* Complex value styles. */
-        span.empty_container::before {
+        span.empty-container::before {
             content: '(empty)';
             font-style: italic;
             margin-left: 0.5em;
@@ -1056,32 +1208,31 @@ class HtmlTreeView(HtmlView):
         """
     )
 
-  @HtmlView.extension_method('_html_tree_view_tooltip')
   def tooltip(
       self,
       value: Any,
       *,
-      name: Optional[str] = None,
       parent: Any = None,
       root_path: Optional[KeyPath] = None,
-      content: Union[str, Html, None] = HtmlView.PresetArgValue(None),
-      **kwargs
+      css_classes: Optional[Sequence[str]] = None,
+      content: Union[str, Html, None] = None,
+      **kwargs,
   ) -> Html:
     """Renders a tooltip for the value.
 
     Args:
       value: The value to render.
-      name: The name of the value.
-      parent: The parent value of the key-value pairs.
+      parent: The parent of the value.
       root_path: The root path of the value.
-      content: The content of the tooltip. If None, the formatted value will be
-        used as the content.
-      **kwargs: Additional keyword arguments passed from `pg.to_html`.
+      css_classes: CSS classes to add to the HTML element.
+      content: The content to render. If None, the value will be rendered.
+      **kwargs: Additional keyword arguments passed from the user that 
+        will be ignored.
 
     Returns:
       The rendered HTML as the tooltip of the value.
     """
-    del name, parent
+    del parent, kwargs
     if content is None:
       content = Html.escape(
           object_utils.format(
@@ -1092,13 +1243,15 @@ class HtmlTreeView(HtmlView):
               python_format=True,
               max_bytes_len=64,
               max_str_len=256,
-              **kwargs
           )
       )
     return Html.element(
         'span',
         [content],
-        css_class=['tooltip', self.css_class_name(value)],
+        css_classes=[
+            'tooltip',
+            css_classes,
+        ],
     ).add_style(
         """
         /* Tooltip styles. */
@@ -1119,23 +1272,220 @@ class HtmlTreeView(HtmlView):
   @staticmethod
   def css_class_name(value: Any) -> Optional[str]:
     """Returns the CSS class name for the value."""
-    if isinstance(value, HtmlTreeView.Extension):
-      return Html.concate(value._html_element_class())  # pylint: disable=protected-access
-    value = value if inspect.isclass(value) else type(value)
-    return object_utils.camel_to_snake(value.__name__, '-')
+    # if isinstance(value, HtmlTreeView.Extension):
+    #   return Html.concate(value._html_element_class())  # pylint: disable=protected-access
+    if inspect.isclass(value):
+      class_name = f'{value.__name__}-class'
+    else:
+      class_name = type(value).__name__
+    return object_utils.camel_to_snake(class_name, '-')
 
   @staticmethod
-  def max_collapse_level(
-      original_level: int | None,
-      subtree_uncollapse_level: int | None,
-      root_path: KeyPath
-  ) -> int | None:
-    """Consolidates the collapse level."""
-    if original_level is None or subtree_uncollapse_level is None:
-      return None
-    return max(original_level, root_path.depth + subtree_uncollapse_level)
+  def init_uncollapse(
+      uncollapse: Union[Iterable[Union[KeyPath, str]], base.NodeFilter, None],
+  ) -> Union[KeyPathSet, base.NodeFilter]:
+    """Initializes the uncollapse argument."""
+    if uncollapse is None:
+      return KeyPathSet()
+    elif callable(uncollapse):
+      return uncollapse
+    else:
+      return KeyPathSet.from_value(uncollapse, include_intermediate=True)
+
+  @staticmethod
+  def get_child_kwargs(
+      call_kwargs: Dict[str, Any],
+      child_config: Dict[str, Any],
+      child_key: Optional[str],
+      root_path: KeyPath,
+  ) -> Dict[str, Any]:
+    """Enter the child config for a child key."""
+    if not child_config:
+      return call_kwargs
 
+    child_kwargs = child_config.get(
+        child_key, child_config.get('__default__', None)
+    )
+    if not child_kwargs:
+      return call_kwargs
+
+    return HtmlTreeView.get_kwargs(
+        call_kwargs, child_kwargs, root_path + child_key,
+    )
 
-_REGEX_CAMEL_TO_SNAKE = re.compile(r'(?<!^)(?=[A-Z])')
+  # pytype: disable=annotation-type-mismatch
+  @staticmethod
+  def get_passthrough_kwargs(
+      *,
+      enable_summary: Optional[bool] = object_utils.MISSING_VALUE,
+      enable_summary_for_str: bool = object_utils.MISSING_VALUE,
+      max_summary_len_for_str: int = object_utils.MISSING_VALUE,
+      enable_summary_tooltip: bool = object_utils.MISSING_VALUE,
+      key_style: Union[
+          Literal['label', 'summary'],
+          Callable[[KeyPath, Any, Any], Literal['label', 'summary']]
+      ] = object_utils.MISSING_VALUE,
+      key_color: Union[
+          Tuple[Optional[str], Optional[str]],
+          Callable[[KeyPath, Any, Any], Tuple[Optional[str], Optional[str]]]
+      ] = object_utils.MISSING_VALUE,
+      include_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = object_utils.MISSING_VALUE,
+      exclude_keys: Union[
+          Iterable[Union[int, str]],
+          Callable[[KeyPath, Any, Any], Iterable[Union[int, str]]],
+          None
+      ] = object_utils.MISSING_VALUE,
+      enable_key_tooltip: bool = object_utils.MISSING_VALUE,
+      uncollapse: Union[
+          KeyPathSet, base.NodeFilter, None
+      ] = object_utils.MISSING_VALUE,
+      extra_flags: Optional[Dict[str, Any]] = object_utils.MISSING_VALUE,
+      highlight: Optional[base.NodeFilter] = object_utils.MISSING_VALUE,
+      lowlight: Optional[base.NodeFilter] = object_utils.MISSING_VALUE,
+      debug: bool = object_utils.MISSING_VALUE,
+      remove: Optional[Iterable[str]] = None,
+      **kwargs,
+  ):
+  # pytype: enable=annotation-type-mismatch
+    """Gets the rendering arguments to pass through to the child nodes."""
+    del kwargs
+    passthrough_kwargs = dict(
+        enable_summary=enable_summary,
+        enable_summary_for_str=enable_summary_for_str,
+        max_summary_len_for_str=max_summary_len_for_str,
+        enable_summary_tooltip=enable_summary_tooltip,
+        enable_key_tooltip=enable_key_tooltip,
+        key_style=key_style,
+        key_color=key_color,
+        include_keys=(
+            include_keys if callable(include_keys)
+            else object_utils.MISSING_VALUE
+        ),
+        exclude_keys=(
+            exclude_keys if callable(exclude_keys)
+            else object_utils.MISSING_VALUE
+        ),
+        uncollapse=uncollapse,
+        highlight=highlight,
+        lowlight=lowlight,
+        extra_flags=extra_flags,
+        debug=debug
+    )
+    # Filter out missing values.
+    passthrough_kwargs = {
+        k: v for k, v in passthrough_kwargs.items()
+        if v is not object_utils.MISSING_VALUE
+    }
+    if remove:
+      return {
+          k: v for k, v in passthrough_kwargs.items()
+          if k not in remove
+      }
+    return passthrough_kwargs
+
+  @staticmethod
+  def get_collapse_level(
+      original_level: Union[None, int, Tuple[Optional[int], int]],
+      overriden_level: Union[None, int, Tuple[Optional[int], int]],
+      ) -> Optional[int]:
+    """Gets the collapse level for a child node."""
+    original_offset, overriden_offset = 0, 0
+    if isinstance(original_level, tuple):
+      original_level, original_offset = original_level
+    if isinstance(overriden_level, tuple):
+      overriden_level, overriden_offset = overriden_level
+
+    if original_level is None:
+      return original_level
+    elif overriden_level is None:
+      return overriden_level
+    else:
+      return max(
+          original_level + original_offset,
+          overriden_level + overriden_offset
+      )
+
+  @staticmethod
+  def get_kwargs(
+      call_kwargs: Dict[str, Any],
+      overriden_kwargs: Dict[str, Any],
+      root_path: Optional[KeyPath] = None,
+  ) -> Dict[str, Any]:
+    """Override render arguments."""
+    # Select child config to override.
+    if not overriden_kwargs:
+      return call_kwargs
+
+    call_kwargs = call_kwargs.copy()
+    overriden_kwargs = overriden_kwargs.copy()
+
+    # Override collapse_level.
+    if 'collapse_level' in call_kwargs or 'collapse_level' in overriden_kwargs:
+      call_kwargs['collapse_level'] = HtmlTreeView.get_collapse_level(
+          call_kwargs.pop('collapse_level', 1),
+          overriden_kwargs.pop('collapse_level', 0)
+      )
+
+    # Override uncollapse.
+    if 'uncollapse' in call_kwargs or 'uncollapse' in overriden_kwargs:
+      uncollapse = KeyPathSet.from_value(
+          call_kwargs.pop('uncollapse', None) or []
+      )
+      child_uncollapse = KeyPathSet.from_value(
+          overriden_kwargs.pop('uncollapse', None) or []
+      )
+      call_kwargs['uncollapse'] = HtmlTreeView.merge_uncollapse(
+          uncollapse, child_uncollapse, root_path
+      )
+
+    # Deep hierarchy merge.
+    return object_utils.merge_tree(call_kwargs, overriden_kwargs)
+
+  @staticmethod
+  def merge_uncollapse(
+      uncollapse: Union[KeyPathSet, base.NodeFilter, None],
+      child_uncollapse: Optional[KeyPathSet],
+      child_path: Optional[KeyPath] = None,
+  ) -> Union[KeyPathSet, base.NodeFilter]:
+    """Merge uncollapse paths."""
+    if not uncollapse and not child_uncollapse:
+      return KeyPathSet()
+
+    if callable(uncollapse) or not child_uncollapse:
+      assert uncollapse is not None
+      return uncollapse
+
+    assert isinstance(uncollapse, KeyPathSet), uncollapse
+    assert isinstance(child_uncollapse, KeyPathSet), child_uncollapse
+    if child_path:
+      child_uncollapse = child_uncollapse.copy()
+      child_uncollapse.rebase(child_path)
+    uncollapse.update(child_uncollapse)
+    return uncollapse
+
+  @staticmethod
+  def get_color(
+      color: Union[
+          Tuple[str, str],
+          Callable[
+              [KeyPath, Any, Any],
+              Tuple[Optional[str], Optional[str]]
+          ],
+          None
+      ],
+      root_path,
+      value,
+      parent
+  ) -> Tuple[Optional[str], Optional[str]]:
+    if callable(color):
+      return color(root_path, value, parent)
+    if color is None:
+      return (None, None)
+    assert isinstance(color, tuple) and len(color) == 2, color
+    return color
 
 # pytype: enable=annotation-type-mismatch