Allow set by tuple

Closes: #48
APN-Pucky · Nov 16, 2024 · 713ee0b · 713ee0b
1 parent ef6224d
commit 713ee0b
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 11 deletions.
diff --git a/src/babyyoda/histo1d.py b/src/babyyoda/histo1d.py
@@ -19,11 +19,24 @@ def set_bin1d(target: Any, source: Any) -> None:
     # self.d_xmin = bin.xMin()
     # self.d_xmax = bin.xMax()
     if hasattr(target, "set"):
-        target.set(
-            source.numEntries(),
-            [source.sumW(), source.sumWX()],
-            [source.sumW2(), source.sumWX2()],
-        )
+        if (
+            hasattr(source, "sumW")
+            and hasattr(source, "sumWX")
+            and hasattr(source, "sumW2")
+            and hasattr(source, "sumWX2")
+            and hasattr(source, "numEntries")
+        ):
+            target.set(
+                source.numEntries(),
+                [source.sumW(), source.sumWX()],
+                [source.sumW2(), source.sumWX2()],
+            )
+        # else if tuple with 3 elements
+        elif len(source) == 3:
+            target.set(source[0], source[1], source[2])
+        else:
+            err = "Invalid argument type"
+            raise NotImplementedError(err)
     else:
         err = "YODA1 backend can not set bin values"
         raise NotImplementedError(err)

diff --git a/src/babyyoda/histo2d.py b/src/babyyoda/histo2d.py
@@ -27,12 +27,26 @@ def set_bin2d(target: Any, source: Any) -> None:
     # self.d_xmin = bin.xMin()
     # self.d_xmax = bin.xMax()
     if hasattr(target, "set"):
-        target.set(
-            source.numEntries(),
-            [source.sumW(), source.sumWX(), source.sumWY()],
-            [source.sumW2(), source.sumWX2(), source.sumWY2()],
-            [source.crossTerm(0, 1)],
-        )
+        if (
+            hasattr(source, "sumW")
+            and hasattr(source, "sumWX")
+            and hasattr(source, "sumWY")
+            and hasattr(source, "sumW2")
+            and hasattr(source, "sumWX2")
+            and hasattr(source, "sumWY2")
+            and hasattr(source, "crossTerm")
+        ):
+            target.set(
+                source.numEntries(),
+                [source.sumW(), source.sumWX(), source.sumWY()],
+                [source.sumW2(), source.sumWX2(), source.sumWY2()],
+                [source.crossTerm(0, 1)],
+            )
+        elif len(source) == 3:
+            target.set(source[0], source[1], source[2])
+        else:
+            err = "Invalid argument type"
+            raise NotImplementedError(err)
     else:
         err = "YODA1 backend can not set bin values"
         raise NotImplementedError(err)