Merge pull request #1884 from braingram/info_v2

Add `Converter.to_info` to customize `info` and `search`

asdf/_asdf.py

@@ -1412,6 +1412,7 @@ def schema_info(self, key="description", path=None, preserve_list=True, refresh_
             self.tree,
             preserve_list=preserve_list,
             refresh_extension_manager=refresh_extension_manager,
+            extension_manager=self.extension_manager,
         )
 
     def info(
@@ -1452,6 +1453,7 @@ def info(
             show_values=show_values,
             identifier="root",
             refresh_extension_manager=refresh_extension_manager,
+            extension_manager=self.extension_manager,
         )
         print("\n".join(lines))
 

asdf/_convenience.py

@@ -7,7 +7,7 @@
 from contextlib import contextmanager
 
 from ._asdf import AsdfFile, open_asdf
-from ._display import DEFAULT_MAX_COLS, DEFAULT_MAX_ROWS, DEFAULT_SHOW_VALUES, render_tree
+from ._display import DEFAULT_MAX_COLS, DEFAULT_MAX_ROWS, DEFAULT_SHOW_VALUES
 
 __all__ = ["info"]
 
@@ -41,18 +41,17 @@ def info(node_or_path, max_rows=DEFAULT_MAX_ROWS, max_cols=DEFAULT_MAX_COLS, sho
         the rendered tree.
     """
     with _manage_node(node_or_path) as node:
-        lines = render_tree(node, max_rows=max_rows, max_cols=max_cols, show_values=show_values, identifier="root")
-        print("\n".join(lines))
+        node.info(max_rows=max_rows, max_cols=max_cols, show_values=show_values)
 
 
 @contextmanager
 def _manage_node(node_or_path):
     if isinstance(node_or_path, (str, pathlib.Path)):
         with open_asdf(node_or_path) as af:
-            yield af.tree
+            yield af
 
     elif isinstance(node_or_path, AsdfFile):
-        yield node_or_path.tree
+        yield node_or_path
 
     else:
-        yield node_or_path
+        yield AsdfFile(node_or_path)

asdf/_core/_converters/ndarray.py

@@ -190,3 +190,6 @@ def data_callback(_attr=None, _ref=weakref.ref(ctx._blocks)):
 
         msg = "Invalid ndarray description."
         raise TypeError(msg)
+
+    def to_info(self, obj):
+        return {"shape": obj.shape, "dtype": obj.dtype}

asdf/_display.py

@@ -8,28 +8,22 @@
 return what it thinks is suitable for display.
 """
 
-import numpy as np
+import sys
 
 from ._node_info import create_tree
-from .tags.core.ndarray import NDArrayType
 
 __all__ = [
     "DEFAULT_MAX_ROWS",
     "DEFAULT_MAX_COLS",
     "DEFAULT_SHOW_VALUES",
     "render_tree",
-    "format_bold",
-    "format_faint",
-    "format_italic",
 ]
 
 
 DEFAULT_MAX_ROWS = 24
 DEFAULT_MAX_COLS = 120
 DEFAULT_SHOW_VALUES = True
 
-EXTENSION_MANAGER = None
-
 
 def render_tree(
     node,
@@ -39,6 +33,7 @@ def render_tree(
     filters=None,
     identifier="root",
     refresh_extension_manager=False,
+    extension_manager=None,
 ):
     """
     Render a tree as text with indents showing depth.
@@ -49,6 +44,7 @@ def render_tree(
         identifier=identifier,
         filters=[] if filters is None else filters,
         refresh_extension_manager=refresh_extension_manager,
+        extension_manager=extension_manager,
     )
     if info is None:
         return []
@@ -61,31 +57,6 @@ def render_tree(
     return renderer.render(info)
 
 
-def format_bold(value):
-    """
-    Wrap the input value in the ANSI escape sequence for increased intensity.
-    """
-    return _format_code(value, 1)
-
-
-def format_faint(value):
-    """
-    Wrap the input value in the ANSI escape sequence for decreased intensity.
-    """
-    return _format_code(value, 2)
-
-
-def format_italic(value):
-    """
-    Wrap the input value in the ANSI escape sequence for italic.
-    """
-    return _format_code(value, 3)
-
-
-def _format_code(value, code):
-    return f"\x1B[{code}m{value}\x1B[0m"
-
-
 class _TreeRenderer:
     """
     Render a _NodeInfo tree with indent showing depth.
@@ -95,14 +66,38 @@ def __init__(self, max_rows, max_cols, show_values):
         self._max_rows = max_rows
         self._max_cols = max_cols
         self._show_values = show_values
+        self._isatty = hasattr(sys.stdout, "isatty") and sys.stdout.isatty()
+
+    def format_bold(self, value):
+        """
+        Wrap the input value in the ANSI escape sequence for increased intensity.
+        """
+        return self._format_code(value, 1)
+
+    def format_faint(self, value):
+        """
+        Wrap the input value in the ANSI escape sequence for decreased intensity.
+        """
+        return self._format_code(value, 2)
+
+    def format_italic(self, value):
+        """
+        Wrap the input value in the ANSI escape sequence for italic.
+        """
+        return self._format_code(value, 3)
+
+    def _format_code(self, value, code):
+        if not self._isatty:
+            return f"{value}"
+        return f"\x1B[{code}m{value}\x1B[0m"
 
     def render(self, info):
         self._mark_visible(info)
 
         lines, elided = self._render(info, set(), True)
 
         if elided:
-            lines.append(format_faint(format_italic("Some nodes not shown.")))
+            lines.append(self.format_faint(self.format_italic("Some nodes not shown.")))
 
         return lines
 
@@ -220,7 +215,7 @@ def _render(self, info, active_depths, is_tail):
         if num_visible_children > 0 and num_visible_children != len(info.children):
             hidden_count = len(info.children) - num_visible_children
             prefix = self._make_prefix(info.depth + 1, active_depths, True)
-            message = format_faint(format_italic(str(hidden_count) + " not shown"))
+            message = self.format_faint(self.format_italic(str(hidden_count) + " not shown"))
             lines.append(f"{prefix}{message}")
 
         return lines, elided
@@ -230,32 +225,29 @@ def _render_node(self, info, active_depths, is_tail):
         value = self._render_node_value(info)
 
         line = (
-            f"{prefix}[{format_bold(info.identifier)}] {value}"
+            f"{prefix}[{self.format_bold(info.identifier)}] {value}"
             if isinstance(info.parent_node, (list, tuple))
-            else f"{prefix}{format_bold(info.identifier)} {value}"
+            else f"{prefix}{self.format_bold(info.identifier)} {value}"
         )
 
         if info.info is not None:
-            line = line + format_faint(format_italic(" # " + info.info))
+            line = line + self.format_faint(self.format_italic(" # " + info.info))
         visible_children = info.visible_children
         if len(visible_children) == 0 and len(info.children) > 0:
-            line = line + format_italic(" ...")
+            line = line + self.format_italic(" ...")
 
         if info.recursive:
-            line = line + " " + format_faint(format_italic("(recursive reference)"))
+            line = line + " " + self.format_faint(self.format_italic("(recursive reference)"))
 
         if self._max_cols is not None and len(line) > self._max_cols:
             message = " (truncated)"
-            line = line[0 : (self._max_cols - len(message))] + format_faint(format_italic(message))
+            line = line[0 : (self._max_cols - len(message))] + self.format_faint(self.format_italic(message))
 
         return line
 
     def _render_node_value(self, info):
         rendered_type = type(info.node).__name__
 
-        if isinstance(info.node, (NDArrayType, np.ndarray)):
-            return f"({rendered_type}): shape={info.node.shape}, dtype={info.node.dtype.name}"
-
         if not info.children and self._show_values:
             try:
                 s = f"{info.node}"
@@ -287,4 +279,4 @@ def _make_prefix(self, depth, active_depths, is_tail):
 
         prefix = prefix + "└─" if is_tail else prefix + "├─"
 
-        return format_faint(prefix)
+        return self.format_faint(prefix)

asdf/_node_info.py

@@ -2,7 +2,7 @@
 from collections import namedtuple
 
 from .schema import load_schema
-from .treeutil import get_children
+from .treeutil import get_children, is_container
 
 
 def _filter_tree(info, filters):
@@ -138,7 +138,7 @@ def _get_schema_key(schema, key):
     return None
 
 
-def create_tree(key, node, identifier="root", filters=None, refresh_extension_manager=False):
+def create_tree(key, node, identifier="root", filters=None, refresh_extension_manager=False, extension_manager=None):
     """
     Create a `NodeSchemaInfo` tree which can be filtered from a base node.
 
@@ -164,6 +164,7 @@ def create_tree(key, node, identifier="root", filters=None, refresh_extension_ma
         identifier,
         node,
         refresh_extension_manager=refresh_extension_manager,
+        extension_manager=extension_manager,
     )
 
     if len(filters) > 0 and not _filter_tree(schema_info, filters):
@@ -180,6 +181,7 @@ def collect_schema_info(
     filters=None,
     preserve_list=True,
     refresh_extension_manager=False,
+    extension_manager=None,
 ):
     """
     Collect from the underlying schemas any of the info stored under key, relative to the path
@@ -209,6 +211,7 @@ def collect_schema_info(
         identifier=identifier,
         filters=[] if filters is None else filters,
         refresh_extension_manager=refresh_extension_manager,
+        extension_manager=extension_manager,
     )
 
     info = schema_info.collect_info(preserve_list=preserve_list)
@@ -235,6 +238,15 @@ def _get_extension_manager(refresh_extension_manager):
     return af.extension_manager
 
 
+def _make_traversable(node, extension_manager):
+    if hasattr(node, "__asdf_traverse__"):
+        return node.__asdf_traverse__(), True, False
+    node_type = type(node)
+    if not extension_manager.handles_type(node_type):
+        return node, False, False
+    return extension_manager.get_converter_for_type(node_type).to_info(node), False, True
+
+
 _SchemaInfo = namedtuple("SchemaInfo", ["info", "value"])
 
 
@@ -299,7 +311,7 @@ class NodeSchemaInfo:
         The portion of the underlying schema corresponding to the node.
     """
 
-    def __init__(self, key, parent, identifier, node, depth, recursive=False, visible=True):
+    def __init__(self, key, parent, identifier, node, depth, recursive=False, visible=True, extension_manager=None):
         self.key = key
         self.parent = parent
         self.identifier = identifier
@@ -309,15 +321,7 @@ def __init__(self, key, parent, identifier, node, depth, recursive=False, visibl
         self.visible = visible
         self.children = []
         self.schema = None
-
-    @classmethod
-    def traversable(cls, node):
-        """
-        This method determines if the node is an instance of a class that
-        supports introspection by the info machinery. This determined by
-        the presence of a __asdf_traverse__ method.
-        """
-        return hasattr(node, "__asdf_traverse__")
+        self.extension_manager = extension_manager or _get_extension_manager()
 
     @property
     def visible_children(self):
@@ -354,13 +358,15 @@ def set_schema_from_node(self, node, extension_manager):
         self.schema = schema
 
     @classmethod
-    def from_root_node(cls, key, root_identifier, root_node, schema=None, refresh_extension_manager=False):
+    def from_root_node(
+        cls, key, root_identifier, root_node, schema=None, refresh_extension_manager=False, extension_manager=None
+    ):
         """
         Build a NodeSchemaInfo tree from the given ASDF root node.
         Intentionally processes the tree in breadth-first order so that recursively
         referenced nodes are displayed at their shallowest reference point.
         """
-        extension_manager = _get_extension_manager(refresh_extension_manager)
+        extension_manager = extension_manager or _get_extension_manager(refresh_extension_manager)
 
         current_nodes = [(None, root_identifier, root_node)]
         seen = set()
@@ -370,26 +376,50 @@ def from_root_node(cls, key, root_identifier, root_node, schema=None, refresh_ex
             next_nodes = []
 
             for parent, identifier, node in current_nodes:
-                if (isinstance(node, (dict, tuple)) or cls.traversable(node)) and id(node) in seen:
-                    info = NodeSchemaInfo(key, parent, identifier, node, current_depth, recursive=True)
+                # node is the item in the tree
+                # We might sometimes not want to use that node directly
+                # but instead using a different node for traversal.
+                t_node, traversable, from_converter = _make_traversable(node, extension_manager)
+                if (is_container(node) or traversable) and id(node) in seen:
+                    info = NodeSchemaInfo(
+                        key,
+                        parent,
+                        identifier,
+                        node,
+                        current_depth,
+                        recursive=True,
+                        extension_manager=extension_manager,
+                    )
                     parent.children.append(info)
 
                 else:
-                    info = NodeSchemaInfo(key, parent, identifier, node, current_depth)
+                    info = NodeSchemaInfo(
+                        key, parent, identifier, node, current_depth, extension_manager=extension_manager
+                    )
 
+                    # If this is the first node keep a reference so we can return it.
                     if root_info is None:
                         root_info = info
 
+                    if parent is None:
+                        info.schema = schema
+
                     if parent is not None:
-                        if parent.schema is not None and not cls.traversable(node):
+                        if parent.schema is not None:
+                            # descend within the schema of the parent
                             info.set_schema_for_property(parent, identifier)
 
+                        # track that this node is a child of the parent
                         parent.children.append(info)
 
+                    # Track which nodes have been seen to avoid an infinite
+                    # loop and to find recursive references
+                    # This is tree wide but should be per-branch.
                     seen.add(id(node))
 
-                    if cls.traversable(node):
-                        t_node = node.__asdf_traverse__()
+                    # if the node has __asdf_traverse__ and a _tag attribute
+                    # that is a valid tag, load it's schema
+                    if traversable:
                         if hasattr(node, "_tag") and isinstance(node._tag, str):
                             try:
                                 info.set_schema_from_node(node, extension_manager)
@@ -400,12 +430,7 @@ def from_root_node(cls, key, root_identifier, root_node, schema=None, refresh_ex
                                 # be using _tag for a non-ASDF purpose.
                                 pass
 
-                    else:
-                        t_node = node
-
-                    if parent is None:
-                        info.schema = schema
-
+                    # add children to queue
                     for child_identifier, child_node in get_children(t_node):
                         next_nodes.append((info, child_identifier, child_node))