feat: allow library to be specified in writer_info removal (#1042)

* feat: allow library to be specified in writer_info removal

Signed-off-by: Henry Schreiner <henryschreineriii@gmail.com>

* style: pre-commit fixes

---------

Signed-off-by: Henry Schreiner <henryschreineriii@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: henryiii

src/boost_histogram/serialization/__init__.py

@@ -46,17 +46,25 @@ def from_uhi(data: dict[str, Any], /) -> histogram.Histogram:
 T = TypeVar("T", bound="dict[str, Any]")
 
 
-def remove_writer_info(obj: T) -> T:
-    """Removes all boost-histogram writer_info from a histogram dict, axes dict, or storage dict. Makes copies where required, and the outer dictionary is always copied."""
+def remove_writer_info(obj: T, /, *, library: str | None = "boost-histogram") -> T:
+    """
+    Removes all ``writer_info`` for a library from a histogram dict, axes dict,
+    or storage dict. Makes copies where required, and the outer dictionary is
+    always copied.
+
+    Specify a library name, or ``None`` to remove all.
+    """
 
     obj = copy.copy(obj)
-    if "boost-histogram" in obj.get("writer_info", {}):
+    if library is None:
+        obj.pop("writer_info")
+    elif library in obj.get("writer_info", {}):
         obj["writer_info"] = copy.copy(obj["writer_info"])
-        del obj["writer_info"]["boost-histogram"]
+        del obj["writer_info"][library]
 
     if "axes" in obj:
-        obj["axes"] = [remove_writer_info(ax) for ax in obj["axes"]]
+        obj["axes"] = [remove_writer_info(ax, library=library) for ax in obj["axes"]]
     if "storage" in obj:
-        obj["storage"] = remove_writer_info(obj["storage"])
+        obj["storage"] = remove_writer_info(obj["storage"], library=library)
 
     return obj

tests/test_serialization_uhi.py

@@ -221,15 +221,16 @@ def test_round_trip_native() -> None:
     assert h2.storage_type is bh.storage.AtomicInt64
 
 
-def test_round_trip_clean() -> None:
+@pytest.mark.parametrize("remove", ["boost-histogram", None])
+def test_round_trip_clean(remove: str | None) -> None:
     h = bh.Histogram(
         bh.axis.Integer(0, 10),
         storage=bh.storage.AtomicInt64(),
     )
     h.fill([-1, 0, 0, 1, 20, 20, 20])
 
     data = to_uhi(h)
-    data = remove_writer_info(data)
+    data = remove_writer_info(data, library=remove)
     h2 = from_uhi(data)
 
     assert isinstance(h2.axes[0], bh.axis.Regular)
@@ -262,3 +263,25 @@ def test_histogram_metadata() -> None:
         "other": 3,
         "_variance_known": True,
     }
+
+
+def test_remove_writer_info() -> None:
+    d = {
+        "uhi_schema": 1,
+        "writer_info": {"boost-histogram": {"foo": "bar"}, "hist": {"FOO": "BAR"}},
+    }
+
+    assert remove_writer_info(d, library=None) == {"uhi_schema": 1}
+    assert remove_writer_info(d) == {
+        "uhi_schema": 1,
+        "writer_info": {"hist": {"FOO": "BAR"}},
+    }
+    assert remove_writer_info(d, library="boost-histogram") == {
+        "uhi_schema": 1,
+        "writer_info": {"hist": {"FOO": "BAR"}},
+    }
+    assert remove_writer_info(d, library="hist") == {
+        "uhi_schema": 1,
+        "writer_info": {"boost-histogram": {"foo": "bar"}},
+    }
+    assert remove_writer_info(d, library="c") == d