add classmethod and serialization to th2 and th3

Author: siliataider

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_th2.py

@@ -83,6 +83,11 @@ def _FillWithArrayTH2(self, *args):
     # Add UHI plotting features
     pythonization(klass)(_add_plotting_features)
 
+    # Add serialization features
+    from ROOT._pythonization._uhi import _add_serialization_features
+
+    pythonization(klass)(_add_serialization_features)
+
     # Support vectorized Fill
     @pythonization(klass)
     def _enable_numpy_fill(klass):

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_th3.py

@@ -8,9 +8,9 @@
 # For the list of contributors see $ROOTSYS/README/CREDITS.                    #
 ################################################################################
 
-from . import pythonization
 from ROOT._pythonization._memory_utils import inject_constructor_releasing_ownership
 
+from . import pythonization
 
 # The constructors need to be pythonized for each derived class separately:
 _th3_derived_classes_to_pythonize = [
@@ -31,3 +31,8 @@
 
     # Add UHI plotting features
     pythonization(klass)(_add_plotting_features)
+
+    # Add serialization features
+    from ROOT._pythonization._uhi import _add_serialization_features
+
+    pythonization(klass)(_add_serialization_features)

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/__init__.py

@@ -77,17 +77,25 @@ def _TH1_Constructor(self, *args, **kwargs):
     else forward to the original
     """
     if len(args) == 1 and isinstance(args[0], dict):
-        from .serialization import _from_uhi_
+        from .serialization import _get_ctor_args, _set_histogram_storage_from_dict
+
+        # get constructor arguments
+        uhi_dict = args[0]
+        ctor_args = _get_ctor_args(self, uhi_dict)
+        # call original constructor
+        self._original_init_(*ctor_args)
+        # set storage
+        _set_histogram_storage_from_dict(self, uhi_dict["storage"])
 
-        _from_uhi_(self, args[0])
     else:
         self._original_init_(*args, **kwargs)
 
 
 def _add_serialization_features(klass: Any) -> None:
-    from .serialization import _to_uhi_
+    from .serialization import _from_uhi_, _to_uhi_
 
     klass._to_uhi_ = _to_uhi_
+    klass._from_uhi_ = classmethod(_from_uhi_)
 
     klass._original_init_ = klass.__init__
     klass.__init__ = _TH1_Constructor

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi/serialization.py

@@ -92,6 +92,20 @@ def _storage_to_dict(hist: Any) -> dict[str, Any]:
     return storage_dict
 
 
+def _get_ctor_args(self, uhi_dict: dict[str, Any]) -> ROOT.TH1:
+    # rebuild axes
+    axes = uhi_dict["axes"]
+    axes_specs = [_axis_from_dict(axis_dict) for axis_dict in axes]
+
+    # constructor arguments
+    # TODO give a unique name or disable gloval directory?
+    ctor_args = ["h_uhi", "h_uhi"]
+    for axis_spec in axes_specs:
+        ctor_args.extend(axis_spec)
+
+    return ctor_args
+
+
 def _set_histogram_storage_from_dict(hist: Any, storage_dict: dict[str, Any]) -> None:
     """
     Set the histogram storage (values and statistics) from the given storage dictionary.
@@ -109,17 +123,17 @@ def _to_uhi_(self) -> dict[str, Any]:
     }
 
 
-def _from_uhi_(self, uhi_dict: dict[str, Any]) -> ROOT.TH1:
-    # rebuild axes
-    axes = uhi_dict["axes"]
-    axes_specs = [_axis_from_dict(axis_dict) for axis_dict in axes]
+def _from_uhi_(cls, uhi_dict: dict[str, Any]) -> ROOT.TH1:
+    # get constructor arguments
+    ctor_args = _get_ctor_args(cls, uhi_dict)
 
-    # construct the histogram
-    ctor_args = ["h_uhi", "h_uhi"]
-    for axis_spec in axes_specs:
-        ctor_args.extend(axis_spec)
+    # allocate instance
+    self = cls.__new__(cls)
 
+    # call original constructor
     self._original_init_(*ctor_args)
 
     # set storage
     _set_histogram_storage_from_dict(self, uhi_dict["storage"])
+
+    return self