feat: metadata and __slots__ for Histogram (#403)

Author: henryiii

src/boost_histogram/_internal/axis.py

@@ -80,6 +80,10 @@ def metadata(self):
     def metadata(self, value):
         self._ax.metadata = value
 
+    @metadata.deleter
+    def metadata(self):
+        self._ax.metadata = None
+
     @classmethod
     def _convert_cpp(cls, cpp_object):
         nice_ax = cls.__new__(cls)

src/boost_histogram/_internal/hist.py

@@ -84,7 +84,11 @@ def _expand_ellipsis(indexes, rank):
 @set_family(MAIN_FAMILY)
 @set_module("boost_histogram")
 class Histogram(object):
-    @inject_signature("self, *axes, storage=Double()", locals={"Double": Double})
+    __slots__ = ("_hist", "axes", "metadata")
+
+    @inject_signature(
+        "self, *axes, storage=Double(), metadata=None", locals={"Double": Double}
+    )
     def __init__(self, *axes, **kwargs):
         """
         Construct a new histogram.
@@ -99,22 +103,27 @@ def __init__(self, *axes, **kwargs):
             Provide 1 or more axis instances.
         storage : Storage = bh.storage.Double()
             Select a storage to use in the histogram
+        metadata : Any = None
+            Data that is passed along if a new histogram is created
         """
 
         # Allow construction from a raw histogram object (internal)
         if not kwargs and len(axes) == 1 and isinstance(axes[0], _histograms):
             self._hist = axes[0]
+            self.metadata = None
             self.axes = self._generate_axes_()
             return
 
+        # If we construct with another Histogram, support that too
         if not kwargs and len(axes) == 1 and isinstance(axes[0], Histogram):
-            self._hist = copy.copy(axes[0]._hist)
-            self.axes = self._generate_axes_()
+            self.__init__(axes[0]._hist)
+            self.metadata = axes[0].metadata
             return
 
         # Keyword only trick (change when Python2 is dropped)
         with KWArgs(kwargs) as k:
             storage = k.optional("storage", Double())
+            self.metadata = k.optional("metadata")
 
         # Check for missed parenthesis or incorrect types
         if not isinstance(storage, Storage):
@@ -150,6 +159,22 @@ def _generate_axes_(self):
 
         return AxesTuple(self._axis(i) for i in range(self.ndim))
 
+    def _new_hist(self, _hist):
+        """
+        Return a new histogram given a new _hist, copying metadata.
+        """
+
+        other = self.__class__(_hist)
+        other.metadata = self.metadata
+        return other
+
+    @property
+    def ndim(self):
+        """
+        Number of axes (dimensions) of histogram.
+        """
+        return self._hist.rank()
+
     def view(self, flow=False):
         """
         Return a view into the data, optionally with overflow turned on.
@@ -161,7 +186,7 @@ def __array__(self):
 
     def __add__(self, other):
         if hasattr(other, "_hist"):
-            return self.__class__(self._hist.__add__(other._hist))
+            return self._new_hist(self._hist.__add__(other._hist))
         else:
             retval = self.copy()
             retval += other
@@ -188,7 +213,7 @@ def __ne__(self, other):
 
     # If these fail, the underlying object throws the correct error
     def __mul__(self, other):
-        return self.__class__(self._hist.__mul__(other))
+        return self._new_hist(self._hist.__mul__(other))
 
     def __rmul__(self, other):
         return self * other
@@ -203,15 +228,15 @@ def __truediv__(self, other):
             result.__itruediv__(other)
             return result
         else:
-            return self.__class__(self._hist.__truediv__(_hist_or_val(other)))
+            return self._new_hist(self._hist.__truediv__(_hist_or_val(other)))
 
     def __div__(self, other):
         if isinstance(other, Histogram):
             result = self.copy()
             result.__idiv__(other)
             return result
         else:
-            return self.__class__(self._hist.__div__(_hist_or_val(other)))
+            return self._new_hist(self._hist.__div__(_hist_or_val(other)))
 
     def __itruediv__(self, other):
         if isinstance(other, Histogram):
@@ -229,12 +254,6 @@ def __idiv__(self, other):
             self._hist.__idiv__(_hist_or_val(other))
         return self
 
-    def __copy__(self):
-        other = self.__class__.__new__(self.__class__)
-        other._hist = copy.copy(self._hist)
-        other.axes = other._generate_axes_()
-        return other
-
     # TODO: Marked as too complex by flake8. Should be factored out a bit.
     @inject_signature("self, *args, weight=None, sample=None, threads=None")
     def fill(self, *args, **kwargs):  # noqa: C901
@@ -347,21 +366,26 @@ def _storage_type(self):
         return cast(self, self._hist._storage_type, Storage)
 
     def _reduce(self, *args):
-        return self.__class__(self._hist.reduce(*args))
+        return self._new_hist(self._hist.reduce(*args))
+
+    def __copy__(self):
+        other = self._new_hist(copy.copy(self._hist))
+        return other
 
     def __deepcopy__(self, memo):
         other = self.__class__.__new__(self.__class__)
         other._hist = copy.deepcopy(self._hist, memo)
+        other.metadata = copy.deepcopy(self.metadata, memo)
         other.axes = other._generate_axes_()
         return other
 
     def __getstate__(self):
-        state = self.__dict__.copy()
-        del state["axes"]  # Don't save the cashe
+        state = {"_hist": self._hist, "metadata": self.metadata}
         return state
 
     def __setstate__(self, state):
-        self.__dict__.update(state)
+        self._hist = state["_hist"]
+        self.metadata = state["metadata"]
         self.axes = self._generate_axes_()
 
     def __repr__(self):
@@ -489,13 +513,6 @@ def rank(self):
         warnings.warn(msg, FutureWarning)
         return self._hist.rank()
 
-    @property
-    def ndim(self):
-        """
-        Number of axes (dimensions) of histogram.
-        """
-        return self._hist.rank()
-
     @property
     def size(self):
         """
@@ -579,12 +596,12 @@ def __getitem__(self, index):  # noqa: C901
         reduced = self._hist.reduce(*slices)
 
         if not integrations:
-            return self.__class__(reduced)
+            return self._new_hist(reduced)
         else:
             projections = [i for i in range(self.ndim) if i not in integrations]
 
             return (
-                self.__class__(reduced.project(*projections))
+                self._new_hist(reduced.project(*projections))
                 if projections
                 else reduced.sum(flow=True)
             )
@@ -701,4 +718,4 @@ def project(self, *args):
         those axes only. Flow bins are used if available.
         """
 
-        return self.__class__(self._hist.project(*args))
+        return self._new_hist(self._hist.project(*args))

src/boost_histogram/_internal/kwargs.py

@@ -15,6 +15,9 @@ def __exit__(self, *args):
         if self.kwargs:
             raise TypeError("Keyword(s) {} not expected".format(", ".join(self.kwargs)))
 
+    def __contains__(self, item):
+        return item in self.kwargs
+
     def required(self, name):
         if name in self.kwargs:
             self.kwargs.pop(name)

tests/test_histogram_indexing.py

@@ -60,6 +60,8 @@ def test_get_1D_histogram():
 def test_get_1D_slice():
     h1 = bh.Histogram(bh.axis.Regular(10, 0, 1))
     h2 = bh.Histogram(bh.axis.Regular(5, 0, 0.5))
+    h1.metadata = {"that": 3}
+
     h1.fill([0.25, 0.25, 0.25, 0.15])
     h2.fill([0.25, 0.25, 0.25, 0.15])
 
@@ -72,6 +74,8 @@ def test_get_1D_slice():
     assert len(h1[2:4].view()) == 2
     assert len(h1[2 : 4 : bh.rebin(2)].view()) == 1
 
+    assert h1[2:4].metadata == {"that": 3}
+
 
 def test_ellipsis():