Merge pull request asdf-format#981 from eslavich/AL-520-search-needs-…

…convenient-way-to-update-node

Add replace method to AsdfSearchResult

CHANGES.rst

@@ -66,6 +66,9 @@
 - Update ``asdftool extensions`` and ``asdftool tags`` to incorporate
   the new extension API. [#988]
 
+- Add ``AsdfSearchResult.replace`` method for assigning new values to
+  search results. [#981]
+
 2.7.5 (unreleased)
 ------------------
 

asdf/search.py

@@ -198,6 +198,26 @@ def _filter(node, identifier):
             show_values=self._show_values
         )
 
+    def replace(self, value):
+        """
+        Assign a new value in place of all leaf nodes in the
+        search results.
+
+        Parameters
+        ----------
+        value : object
+        """
+        results = []
+
+        def _callback(identifiers, parent, node, children):
+            if all(f(node, identifiers[-1]) for f in self._filters):
+                results.append((identifiers[-1], parent))
+
+        _walk_tree_breadth_first(self._identifiers, self._node, _callback)
+
+        for identifier, parent in results:
+            parent[identifier] = value
+
     @property
     def node(self):
         """
@@ -247,7 +267,7 @@ def nodes(self):
         """
         results = []
 
-        def _callback(identifiers, node, children):
+        def _callback(identifiers, parent, node, children):
             if all(f(node, identifiers[-1]) for f in self._filters):
                 results.append(node)
 
@@ -266,7 +286,7 @@ def paths(self):
         """
         results = []
 
-        def _callback(identifiers, node, children):
+        def _callback(identifiers, parent, node, children):
             if all(f(node, identifiers[-1]) for f in self._filters):
                 results.append(_build_path(identifiers))
 
@@ -310,18 +330,18 @@ def _walk_tree_breadth_first(root_identifiers, root_node, callback):
     Walk the tree in breadth-first order (useful for prioritizing
     lower-depth nodes).
     """
-    current_nodes = [(root_identifiers, root_node)]
+    current_nodes = [(root_identifiers, None, root_node)]
     seen = set()
     while True:
         next_nodes = []
 
-        for identifiers, node in current_nodes:
+        for identifiers, parent, node in current_nodes:
             if (isinstance(node, dict) or isinstance(node, list) or isinstance(node, tuple)) and id(node) in seen:
                 continue
 
             children = get_children(node)
-            callback(identifiers, node, [c for _, c in children])
-            next_nodes.extend([(identifiers + [i], c) for i, c in children])
+            callback(identifiers, parent, node, [c for _, c in children])
+            next_nodes.extend([(identifiers + [i], node, c) for i, c in children])
             seen.add(id(node))
 
         if len(next_nodes) == 0:

asdf/tests/test_search.py

@@ -37,6 +37,8 @@ def test_single_result(asdf_file):
     assert len(result.nodes) == 1
     assert result.node == "hello"
     assert result.path == "root['bar']"
+    result.replace("goodbye")
+    assert asdf_file["bar"] == "goodbye"
 
 
 def test_multiple_results(asdf_file):
@@ -54,6 +56,10 @@ def test_multiple_results(asdf_file):
     with pytest.raises(RuntimeError):
         result.node
 
+    result.replace(54)
+    assert asdf_file["foo"] == 54
+    assert asdf_file["nested"]["foo"] == 54
+
 
 def test_by_key(asdf_file):
     result = asdf_file.search("bar")
@@ -120,12 +126,20 @@ def test_multiple_conditions(asdf_file):
     result = asdf_file.search("foo", value=24)
     assert len(result.nodes) == 1
     assert result.node == 24
+    result.replace(19)
+    assert len(result.nodes) == 0
+    assert asdf_file["foo"] == 42
+    assert asdf_file["nested"]["foo"] == 19
 
 
 def test_chaining(asdf_file):
     result = asdf_file.search("foo").search(value=24)
     assert len(result.nodes) == 1
     assert result.node == 24
+    result.replace(19)
+    assert len(result.nodes) == 0
+    assert asdf_file["foo"] == 42
+    assert asdf_file["nested"]["foo"] == 19
 
 
 def test_index_operator(asdf_file):
@@ -158,6 +172,8 @@ def test_no_results(asdf_file):
     assert "No results found." in repr(result)
     assert result.node is None
     assert result.path is None
+    # Testing no exceptions here:
+    result.replace("foo")
 
 
 def test_recursive_tree():
@@ -171,3 +187,7 @@ def test_recursive_tree():
     result = af.search("bar")
     assert len(result.nodes) == 1
     assert result.node == "baz"
+
+    result.replace("zap")
+    assert af["foo"]["bar"] == "zap"
+    assert af["foo"]["nested"]["bar"] == "zap"

docs/asdf/features.rst

@@ -594,6 +594,14 @@ properties instead:
     >>> af.search('duplicate_key').nodes # doctest: +SKIP
     ["value 1", "value 2"]
 
+To replace matching nodes with a new value, use the `AsdfSearchResult.replace` method:
+
+.. code:: python
+
+    >>> af.search('example').replace('replacement value') # doctest: +SKIP
+    >>> af.search('example').node # doctest: +SKIP
+    'replacement value'
+
 .. currentmodule:: asdf
 
 The first argument to `AsdfFile.search` searches by dict key or list/tuple index.  We can