From 802b7591dbea767c48329170ce573a76acf6290b Mon Sep 17 00:00:00 2001
From: Edward Slavich <ed.slavich@justappraised.com>
Date: Sat, 7 Jan 2023 22:25:10 -0900
Subject: [PATCH] Add tree_mapper argument to asdf.open

---
 CHANGES.rst            |  3 +++
 asdf/asdf.py           | 14 ++++++++++++++
 asdf/tests/test_api.py | 14 ++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/CHANGES.rst b/CHANGES.rst
index ac36e9c15..c6207f0a8 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -6,6 +6,9 @@ The ASDF Standard is at v1.6.0
 - Discard cache of lazy-loaded block data when it is no longer referenced
   by the tree. [#1280]
 
+- Add ``tree_mapper`` argument to ``asdf.open`` to permit arbitrary
+  transformations of the tagged tree prior to conversion.  [#1289]
+
 2.14.3 (2022-12-15)
 -------------------
 
diff --git a/asdf/asdf.py b/asdf/asdf.py
index c65a7aeff..d4fe4f70d 100644
--- a/asdf/asdf.py
+++ b/asdf/asdf.py
@@ -802,6 +802,7 @@ def _open_asdf(
         _force_raw_types=False,
         strict_extension_check=False,
         ignore_missing_extensions=False,
+        tree_mapper=None,
         **kwargs,
     ):
         """Attempt to populate AsdfFile data from file-like object"""
@@ -867,6 +868,9 @@ def _open_asdf(
                 self._blocks.read_internal_blocks(fd, past_magic=True, validate_checksums=validate_checksums)
                 self._blocks.read_block_index(fd, self)
 
+            if tree_mapper is not None:
+                tree = tree_mapper(tree)
+
             tree = reference.find_references(tree, self)
 
             if self.version <= versioning.FILL_DEFAULTS_MAX_VERSION and get_config().legacy_fill_schema_defaults:
@@ -902,6 +906,7 @@ def _open_impl(
         _force_raw_types=False,
         strict_extension_check=False,
         ignore_missing_extensions=False,
+        tree_mapper=None,
         **kwargs,
     ):
         """Attempt to open file-like object as either AsdfFile or AsdfInFits"""
@@ -918,6 +923,7 @@ def _open_impl(
                 _force_raw_types,
                 strict_extension_check,
                 ignore_missing_extensions,
+                tree_mapper=tree_mapper,
                 **kwargs,
             )
         except Exception as e:
@@ -937,6 +943,7 @@ def _open_generic_file(
         _force_raw_types=False,
         strict_extension_check=False,
         ignore_missing_extensions=False,
+        tree_mapper=None,
         **kwargs,
     ):
         """Attempt to open a generic_file instance as either AsdfFile or AsdfInFits"""
@@ -980,6 +987,7 @@ def _open_generic_file(
                 _force_raw_types=_force_raw_types,
                 strict_extension_check=strict_extension_check,
                 ignore_missing_extensions=ignore_missing_extensions,
+                tree_mapper=tree_mapper,
                 **kwargs,
             )
         else:
@@ -1733,6 +1741,7 @@ def open_asdf(
     strict_extension_check=False,
     ignore_missing_extensions=False,
     _compat=False,
+    tree_mapper=None,
     **kwargs,
 ):
     """
@@ -1807,6 +1816,10 @@ def open_asdf(
         and custom schemas.  Recommended unless the file is already known
         to be valid.
 
+    tree_mapper : callable, optional
+        Callable that accepts the tagged ASDF tree before conversion
+        and returns a modified tree.
+
     Returns
     -------
     asdffile : AsdfFile
@@ -1840,6 +1853,7 @@ def open_asdf(
         _force_raw_types=_force_raw_types,
         strict_extension_check=strict_extension_check,
         ignore_missing_extensions=ignore_missing_extensions,
+        tree_mapper=tree_mapper,
         **kwargs,
     )
 
diff --git a/asdf/tests/test_api.py b/asdf/tests/test_api.py
index 2660ffad5..c1303907c 100644
--- a/asdf/tests/test_api.py
+++ b/asdf/tests/test_api.py
@@ -530,3 +530,17 @@ def test_none_values(tmp_path):
     with asdf.open(path) as af:
         assert "foo" in af
         assert af["foo"] is None
+
+
+def test_tree_mapper(tmp_path):
+    path = str(tmp_path / "test.asdf")
+
+    af = asdf.AsdfFile({"foo": "bar"})
+    af.write_to(path)
+
+    def tree_mapper(tree):
+        tree["foo"] = "baz"
+        return tree
+
+    with asdf.open(path, tree_mapper=tree_mapper) as af:
+        assert af["foo"] == "baz"