fix: support __dict__ on histogram like axes (#1041)

* fix: support __dict__ on histogram like axes

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

* fix: ignore None values

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/histogram.py

@@ -226,27 +226,35 @@ def __init_subclass__(cls, *, family: object | None = None) -> None:
         cls._family = family if family is not None else object()
 
     @typing.overload
-    def __init__(self, arg: Histogram, /, *, metadata: Any = ...) -> None: ...
+    def __init__(
+        self, arg: Histogram, /, *, metadata: Any = ..., __dict__: Any = ...
+    ) -> None: ...
 
     @typing.overload
-    def __init__(self, arg: dict[str, Any], /, *, metadata: Any = ...) -> None: ...
+    def __init__(
+        self, arg: dict[str, Any], /, *, metadata: Any = ..., __dict__: Any = ...
+    ) -> None: ...
 
     @typing.overload
-    def __init__(self, arg: CppHistogram, /, *, metadata: Any = ...) -> None: ...
+    def __init__(
+        self, arg: CppHistogram, /, *, metadata: Any = ..., __dict__: Any = ...
+    ) -> None: ...
 
     @typing.overload
     def __init__(
         self,
         *axes: Axis | CppAxis,
         storage: Storage = ...,
         metadata: Any = ...,
+        __dict__: Any = ...,
     ) -> None: ...
 
     def __init__(
         self,
         *axes: Axis | CppAxis | Histogram | CppHistogram | dict[str, Any],
         storage: Storage | None = None,
         metadata: Any = NO_METADATA,
+        __dict__: Any = None,
     ) -> None:
         """
         Construct a new histogram.
@@ -263,33 +271,45 @@ def __init__(
             Select a storage to use in the histogram
         metadata : Any = None
             Data that is passed along if a new histogram is created. No not use
-            in new code; set properties in __dict__ directly instead.
+            in new code; use ``__dict__`` instead.
+        __dict__ : Any = None
+            Better way to set metadata.
         """
         self._variance_known = True
         storage_err_msg = "storage= is not allowed with conversion constructor"
 
+        if metadata is not NO_METADATA and __dict__:
+            msg = (
+                "Can't set both metadata and __dict__. Set the 'metadata' key instead."
+            )
+            raise TypeError(msg)
+        if metadata is not NO_METADATA:
+            __dict__ = {"metadata": metadata}
+        if __dict__ is None:
+            __dict__ = {}
+
         # Allow construction from a raw histogram object (internal)
         if len(axes) == 1 and isinstance(axes[0], tuple(_histograms)):
             if storage is not None:
                 raise TypeError(storage_err_msg)
             cpp_hist: CppHistogram = axes[0]  # type: ignore[assignment]
-            self._from_histogram_cpp(cpp_hist, metadata=metadata)
+            self._from_histogram_cpp(cpp_hist, __dict__=__dict__)
             return
 
         # If we construct with another Histogram as the only positional argument,
         # support that too
         if len(axes) == 1 and isinstance(axes[0], Histogram):
             if storage is not None:
                 raise TypeError(storage_err_msg)
-            self._from_histogram_object(axes[0], metadata=metadata)
+            self._from_histogram_object(axes[0], __dict__=__dict__)
             return
 
         # Support objects that provide a to_boost method, like Uproot
         if len(axes) == 1 and hasattr(axes[0], "_to_boost_histogram_"):
             if storage is not None:
                 raise TypeError(storage_err_msg)
             self._from_histogram_object(
-                axes[0]._to_boost_histogram_(), metadata=metadata
+                axes[0]._to_boost_histogram_(), __dict__=__dict__
             )
             return
 
@@ -298,15 +318,14 @@ def __init__(
             if storage is not None:
                 raise TypeError(storage_err_msg)
             self._from_histogram_object(
-                serialization.from_uhi(axes[0]), metadata=metadata
+                serialization.from_uhi(axes[0]), __dict__=__dict__
             )
             return
 
         if storage is None:
             storage = Double()
 
-        if metadata is not NO_METADATA:
-            self.metadata = metadata
+        self.__dict__.update(__dict__)
 
         # Check for missed parenthesis or incorrect types
         if not isinstance(storage, Storage):
@@ -347,7 +366,7 @@ def _clone(
 
         self = cls.__new__(cls)
         if isinstance(_hist, tuple(_histograms)):
-            self._from_histogram_cpp(_hist, metadata=NO_METADATA)  # type: ignore[arg-type]
+            self._from_histogram_cpp(_hist, __dict__={})  # type: ignore[arg-type]
             if other is not None:
                 return cls._clone(self, other=other, memo=memo)
             return self
@@ -357,18 +376,19 @@ def _clone(
         if other is None:
             other = _hist
 
-        self._from_histogram_object(_hist, metadata=NO_METADATA)
-
         if memo is NOTHING:
-            self.__dict__ = copy.copy(other.__dict__)
+            dict_copy = copy.copy(other.__dict__)
         else:
-            self.__dict__ = copy.deepcopy(other.__dict__, memo)
+            dict_copy = copy.deepcopy(other.__dict__, memo)
+
+        self._from_histogram_object(_hist, __dict__=dict_copy)
 
         for ax in self.axes:
             if memo is NOTHING:
-                ax.__dict__ = copy.copy(ax._ax.raw_metadata)
+                ax._ax.raw_metadata = copy.copy(ax._ax.raw_metadata)
             else:
-                ax.__dict__ = copy.deepcopy(ax._ax.raw_metadata, memo)
+                ax._ax.raw_metadata = copy.deepcopy(ax._ax.raw_metadata, memo)
+            ax.__dict__ = ax._ax.raw_metadata
         return self
 
     def _new_hist(self, _hist: CppHistogram, memo: Any = NOTHING) -> Self:
@@ -377,17 +397,20 @@ def _new_hist(self, _hist: CppHistogram, memo: Any = NOTHING) -> Self:
         """
         return self.__class__._clone(_hist, other=self, memo=memo)
 
-    def _from_histogram_cpp(self, other: CppHistogram, *, metadata: Any) -> None:
+    def _from_histogram_cpp(
+        self, other: CppHistogram, *, __dict__: dict[str, Any]
+    ) -> None:
         """
         Import a Cpp histogram.
         """
         self._variance_known = True
         self._hist = other
-        if metadata is not NO_METADATA:
-            self.metadata = metadata
+        self.__dict__.update(__dict__)
         self.axes = self._generate_axes_()
 
-    def _from_histogram_object(self, other: Histogram, *, metadata: Any) -> None:
+    def _from_histogram_object(
+        self, other: Histogram, *, __dict__: dict[str, Any]
+    ) -> None:
         """
         Convert self into a new histogram object based on another, possibly
         converting from a different subclass.
@@ -396,11 +419,9 @@ def _from_histogram_object(self, other: Histogram, *, metadata: Any) -> None:
         self.__dict__ = copy.copy(other.__dict__)
         self.axes = self._generate_axes_()
         for ax in self.axes:
-            ax.__dict__ = copy.copy(ax._ax.raw_metadata)
-        if metadata is not NO_METADATA:
-            self.metadata = metadata
-        elif "metadata" in other.__dict__:
-            self.metadata = other.metadata
+            ax.__dict__.update(ax._ax.raw_metadata)
+        self.__dict__.update(other.__dict__)
+        self.__dict__.update(__dict__)
 
         # Allow custom behavior on either "from" or "to"
         other._export_bh_(self)

src/boost_histogram/serialization/__init__.py

@@ -6,6 +6,7 @@
 # pylint: disable-next=import-error
 from .. import histogram, version
 from ._axis import _axis_from_dict, _axis_to_dict
+from ._common import serialize_metadata
 from ._storage import _data_from_dict, _storage_from_dict, _storage_to_dict
 
 __all__ = ["from_uhi", "remove_writer_info", "to_uhi"]
@@ -25,7 +26,7 @@ def to_uhi(h: histogram.Histogram, /) -> dict[str, Any]:
         "axes": [_axis_to_dict(axis) for axis in h.axes],
         "storage": _storage_to_dict(h.storage_type(), h.view(flow=True)),
     }
-    data["metadata"] = {k: v for k, v in h.__dict__.items() if not k.startswith("@")}
+    data["metadata"] = serialize_metadata(h.__dict__)
 
     return data
 

src/boost_histogram/serialization/_axis.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 from .. import axis
+from ._common import serialize_metadata
 
 __all__ = ["_axis_from_dict", "_axis_to_dict"]
 
@@ -55,10 +56,8 @@ def _(ax: axis.Regular | axis.Integer, /) -> dict[str, Any]:
 
     if isinstance(ax, axis.Integer):
         data["writer_info"] = {"boost-histogram": {"orig_type": "Integer"}}
-    if ax.metadata is not None:
-        data["metadata"] = {
-            k: v for k, v in ax.metadata.items() if not k.startswith("@")
-        }
+
+    data["metadata"] = serialize_metadata(ax.__dict__)
 
     return data
 
@@ -73,10 +72,7 @@ def _(ax: axis.Variable, /) -> dict[str, Any]:
         "overflow": ax.traits.overflow,
         "circular": ax.traits.circular,
     }
-    if ax.metadata is not None:
-        data["metadata"] = {
-            k: v for k, v in ax.metadata.items() if not k.startswith("@")
-        }
+    data["metadata"] = serialize_metadata(ax.__dict__)
 
     return data
 
@@ -89,10 +85,7 @@ def _(ax: axis.IntCategory, /) -> dict[str, Any]:
         "categories": list(ax),
         "flow": ax.traits.overflow,
     }
-    if ax.metadata is not None:
-        data["metadata"] = {
-            k: v for k, v in ax.metadata.items() if not k.startswith("@")
-        }
+    data["metadata"] = serialize_metadata(ax.__dict__)
 
     return data
 
@@ -105,10 +98,7 @@ def _(ax: axis.StrCategory, /) -> dict[str, Any]:
         "categories": list(ax),
         "flow": ax.traits.overflow,
     }
-    if ax.metadata is not None:
-        data["metadata"] = {
-            k: v for k, v in ax.metadata.items() if not k.startswith("@")
-        }
+    data["metadata"] = serialize_metadata(ax.__dict__)
 
     return data
 
@@ -119,10 +109,7 @@ def _(ax: axis.Boolean, /) -> dict[str, Any]:
     data: dict[str, Any] = {
         "type": "boolean",
     }
-    if ax.metadata is not None:
-        data["metadata"] = {
-            k: v for k, v in ax.metadata.items() if not k.startswith("@")
-        }
+    data["metadata"] = serialize_metadata(ax.__dict__)
 
     return data
 
@@ -139,7 +126,7 @@ def _axis_from_dict(data: dict[str, Any], /) -> axis.Axis:
             underflow=data["underflow"],
             overflow=data["overflow"],
             circular=data["circular"],
-            metadata=data.get("metadata"),
+            __dict__=data.get("metadata"),
         )
 
     hist_type = data["type"]
@@ -151,29 +138,29 @@ def _axis_from_dict(data: dict[str, Any], /) -> axis.Axis:
             underflow=data["underflow"],
             overflow=data["overflow"],
             circular=data["circular"],
-            metadata=data.get("metadata"),
+            __dict__=data.get("metadata"),
         )
     if hist_type == "variable":
         return axis.Variable(
             data["edges"],
             underflow=data["underflow"],
             overflow=data["overflow"],
             circular=data["circular"],
-            metadata=data.get("metadata"),
+            __dict__=data.get("metadata"),
         )
     if hist_type == "category_int":
         return axis.IntCategory(
             data["categories"],
             overflow=data["flow"],
-            metadata=data.get("metadata"),
+            __dict__=data.get("metadata"),
         )
     if hist_type == "category_str":
         return axis.StrCategory(
             data["categories"],
             overflow=data["flow"],
-            metadata=data.get("metadata"),
+            __dict__=data.get("metadata"),
         )
     if hist_type == "boolean":
-        return axis.Boolean(metadata=data.get("metadata"))
+        return axis.Boolean(__dict__=data.get("metadata"))
 
     raise TypeError(f"Unsupported axis type: {hist_type}")

src/boost_histogram/serialization/_common.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def serialize_metadata(value: dict[str, Any], /) -> dict[str, Any]:
+    return {k: v for k, v in value.items() if not k.startswith("@") and v is not None}

tests/test_serialization_uhi.py

@@ -238,7 +238,7 @@ def test_round_trip_clean() -> None:
 
 def test_unserializable_metadata() -> None:
     h = bh.Histogram(
-        bh.axis.Integer(0, 10, metadata={"c": 3, "@d": 4}),
+        bh.axis.Integer(0, 10, __dict__={"c": 3, "@d": 4}),
     )
     h.__dict__["a"] = 1
     h.__dict__["@b"] = 2