Repository: asdf-format/asdf

Issue:
Provide a way of getting schema information given a node

Discussion:

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index efb9fddd..1e3da514 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -3,6 +3,8 @@
 
 - Update citation. [#1184]
 - Add search support to `~asdf.AsdfFile.schema_info`. [#1187]
+- Add `asdf.search.AsdfSearchResult` support for `~asdf.AsdfFile.schema_info` and
+  `~asdf.search.AsdfSearchResult.schema_info` method. [#1197]
 
 2.13.0 (2022-08-19)
 -------------------
diff --git a/asdf/_display.py b/asdf/_display.py
index 173b1871..ea98be66 100644
--- a/asdf/_display.py
+++ b/asdf/_display.py
@@ -9,7 +9,7 @@ return what it thinks is suitable for display.
 """
 import numpy as np
 
-from ._node_info import NodeSchemaInfo
+from ._node_info import create_tree
 from .tags.core.ndarray import NDArrayType
 from .util import is_primitive
 
@@ -43,11 +43,15 @@ def render_tree(
     """
     Render a tree as text with indents showing depth.
     """
-    info = NodeSchemaInfo.from_root_node("title", identifier, node, refresh_extension_manager=refresh_extension_manager)
-
-    if len(filters) > 0:
-        if not _filter_tree(info, filters):
-            return []
+    info = create_tree(
+        key="title",
+        node=node,
+        identifier=identifier,
+        filters=filters,
+        refresh_extension_manager=refresh_extension_manager,
+    )
+    if info is None:
+        return []
 
     renderer = _TreeRenderer(
         max_rows,
@@ -82,20 +86,6 @@ def _format_code(value, code):
     return f"\x1B[{code}m{value}\x1B[0m"
 
 
-def _filter_tree(info, filters):
-    """
-    Remove nodes from the tree that get caught in the filters.
-    Mutates the tree.
-    """
-    filtered_children = []
-    for child in info.children:
-        if _filter_tree(child, filters):
-            filtered_children.append(child)
-    info.children = filtered_children
-
-    return len(info.children) > 0 or all(f(info.node, info.identifier) for f in filters)
-
-
 class _TreeRenderer:
     """
     Render a _NodeInfo tree with indent showing depth.
diff --git a/asdf/_node_info.py b/asdf/_node_info.py
index ce516cad..e989b713 100644
--- a/asdf/_node_info.py
+++ b/asdf/_node_info.py
@@ -5,7 +5,55 @@ from .schema import load_schema
 from .treeutil import get_children
 
 
-def collect_schema_info(key, path, node, identifier="root", preserve_list=True, refresh_extension_manager=False):
+def _filter_tree(info, filters):
+    """
+    Keep nodes in tree which pass the all of the filters.
+
+    Mutates the tree.
+    """
+    filtered_children = []
+    for child in info.children:
+        if _filter_tree(child, filters):
+            filtered_children.append(child)
+    info.children = filtered_children
+
+    return len(info.children) > 0 or all(f(info.node, info.identifier) for f in filters)
+
+
+def create_tree(key, node, identifier="root", filters=[], refresh_extension_manager=False):
+    """
+    Create a `NodeSchemaInfo` tree which can be filtered from a base node.
+
+    Parameters
+    ----------
+    key : str
+        The key to look up.
+    node : dict
+        The asdf tree to search.
+    filters : list of functions
+        A list of functions that take a node and identifier and return True if the node should be included in the tree.
+    preserve_list : bool
+        If True, then lists are preserved. Otherwise, they are turned into dicts.
+    refresh_extension_manager : bool
+        If `True`, refresh the extension manager before looking up the
+        key.  This is useful if you want to make sure that the schema
+        data for a given key is up to date.
+    """
+
+    schema_info = NodeSchemaInfo.from_root_node(
+        key, identifier, node, refresh_extension_manager=refresh_extension_manager
+    )
+
+    if len(filters) > 0:
+        if not _filter_tree(schema_info, filters):
+            return
+
+    return schema_info
+
+
+def collect_schema_info(
+    key, path, node, identifier="root", filters=[], preserve_list=True, refresh_extension_manager=False
+):
     """
     Collect from the underlying schemas any of the info stored under key, relative to the path
 
@@ -18,6 +66,8 @@ def collect_schema_info(key, path, node, identifier="root", preserve_list=True,
         If None return full dictionary.
     node : dict
         The asdf tree to search.
+    filters : list of functions
+        A list of functions that take a node and identifier and return True if the node should be included in the tree.
     preserve_list : bool
         If True, then lists are preserved. Otherwise, they are turned into dicts.
     refresh_extension_manager : bool
@@ -26,8 +76,8 @@ def collect_schema_info(key, path, node, identifier="root", preserve_list=True,
         data for a given key is up to date.
     """
 
-    schema_info = NodeSchemaInfo.from_root_node(
-        key, identifier, node, refresh_extension_manager=refresh_extension_manager
+    schema_info = create_tree(
+        key, node, identifier=identifier, filters=filters, refresh_extension_manager=refresh_extension_manager
     )
 
     info = schema_info.collect_info(preserve_list=preserve_list)
diff --git a/asdf/asdf.py b/asdf/asdf.py
index a12b8242..105fb7ca 100644
--- a/asdf/asdf.py
+++ b/asdf/asdf.py
@@ -1515,8 +1515,9 @@ class AsdfFile:
         key : str
             The key to look up.
             Default: "description"
-        path : str
-            A dot-separated path to the parameter to find the key information on.
+        path : str or asdf.search.AsdfSearchResult
+            A dot-separated path to the parameter to find the key information on or
+            an `asdf.search.AsdfSearchResult` object.
             Default = None (full dictionary).
         preserve_list : bool
             If True, then lists are preserved. Otherwise, they are turned into dicts.
@@ -1526,9 +1527,14 @@ class AsdfFile:
             data for a given key is up to date.
         """
 
-        return node_info.collect_schema_info(
-            key, path, self.tree, preserve_list=preserve_list, refresh_extension_manager=refresh_extension_manager
-        )
+        if isinstance(path, AsdfSearchResult):
+            return path.schema_info(
+                key, preserve_list=preserve_list, refresh_extension_manager=refresh_extension_manager
+            )
+        else:
+            return node_info.collect_schema_info(
+                key, path, self.tree, preserve_list=preserve_list, refresh_extension_manager=refresh_extension_manager
+            )
 
     def info(
         self,
diff --git a/asdf/search.py b/asdf/search.py
index ca066eda..f53129c8 100644
--- a/asdf/search.py
+++ b/asdf/search.py
@@ -7,7 +7,7 @@ import re
 import typing
 
 from ._display import DEFAULT_MAX_COLS, DEFAULT_MAX_ROWS, DEFAULT_SHOW_VALUES, format_faint, format_italic, render_tree
-from ._node_info import NodeSchemaInfo
+from ._node_info import NodeSchemaInfo, collect_schema_info
 from .treeutil import get_children, is_container
 from .util import NotSet
 
@@ -323,6 +323,32 @@ class AsdfSearchResult:
         else:
             return "\n".join(lines)
 
+    def schema_info(self, key="description", preserve_list=True, refresh_extension_manager=False):
+        """
+        Get a nested dictionary of the schema information for a given key, relative to this search result.
+
+        Parameters
+        ----------
+        key : str
+            The key to look up.
+            Default: "description"
+        preserve_list : bool
+            If True, then lists are preserved. Otherwise, they are turned into dicts.
+        refresh_extension_manager : bool
+            If `True`, refresh the extension manager before looking up the
+            key.  This is useful if you want to make sure that the schema
+            data for a given key is up to date.
+        """
+
+        return collect_schema_info(
+            key,
+            None,
+            self._node,
+            filters=self._filters,
+            preserve_list=preserve_list,
+            refresh_extension_manager=refresh_extension_manager,
+        )
+
     def __getitem__(self, key):
         if (
             isinstance(self._node, dict)
diff --git a/docs/asdf/features.rst b/docs/asdf/features.rst
index c619f135..be3cf5f2 100644
--- a/docs/asdf/features.rst
+++ b/docs/asdf/features.rst
@@ -754,8 +754,8 @@ the information from the schema requested together with the value stored in the
 ASDF file itself.
 
 One needs to provide a ``key``, which corresponds the to the keyword the information
-is stored under inside the schema, by default this is ``description``. Moreover,
-one can also provide a ``path`` which is a dot-separated string of the keys in the
+is stored under inside the schema, by default this is ``description``. One can also
+provide a ``path`` in the form of a dot-separated string of the keys in the
 ASDF file that lead to the value(s) of interest. For example:
 
 .. code:: pycon
@@ -763,3 +763,17 @@ ASDF file that lead to the value(s) of interest. For example:
     >>> af.schema_info("archive_catalog", "foo.bar")  # doctest: +SKIP
     {'thing1': {'archive_catalog': 'Thing 1 Archive catalog information'},
      'thing2': {'archive_catalog': 'Thing 2 Archive catalog information'}}
+
+Or one can provide a ``path`` as an `asdf.search.AsdfSearchResult` object:
+
+.. code:: pycon
+
+    >>> af.schema_info("archive_catalog", af.search("bar"))  # doctest: +SKIP
+    {'thing1': {'archive_catalog': 'Thing 1 Archive catalog information'},
+     'thing2': {'archive_catalog': 'Thing 2 Archive catalog information'}}
+
+
+.. note::
+    The there is also the `asdf.search.AsdfSearchResult.schema_info` method,
+    which can be directly called on an `asdf.search.AsdfSearchResult` object.
+    instead of having to pass the search through `AsdfFile.schema_info`.


Test patch:
diff --git a/asdf/tests/test_info.py b/asdf/tests/test_info.py
index bb1c233e..fda42eb2 100644
--- a/asdf/tests/test_info.py
+++ b/asdf/tests/test_info.py
@@ -567,6 +567,17 @@ def test_schema_info_support(tmpdir):
         },
     }
 
+    # Test using a search result
+    search = af.search("clown")
+    assert af.schema_info("description", search, refresh_extension_manager=True) == {
+        "object": {
+            "clown": {
+                "description": ("clown description", "Bozo"),
+            },
+            "description": ("object with info support description", af.tree["object"]),
+        },
+    }
+
 
 def test_info_object_support(capsys, tmpdir):
     manifest_extension(tmpdir)