warn in print() when user is about to compute dask stuff

ikrommyd · Sep 17, 2023 · 3bdff8c · 3bdff8c
1 parent 2fee783
commit 3bdff8c
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/src/coffea/analysis_tools.py b/src/coffea/analysis_tools.py
@@ -582,10 +582,16 @@ def to_npz(self, file, compressed=False, compute=False):
         else:
             return out
 
-    def print(self):
+    def print(self, compute=False):
         """Prints the statistics of the N-1 selection"""
 
-        if self._delayed_mode:
+        if self._delayed_mode and not compute:
+            warnings.warn(
+                "This will compute dask_awkward arrays. If you really want to do this now, call print(compute=True)"
+            )
+            return
+
+        if self._delayed_mode and compute:
             self._nev = list(dask.compute(*self._nev))
 
         nev = self._nev
@@ -830,10 +836,16 @@ def to_npz(self, file, compressed=False, compute=False):
         else:
             return out
 
-    def print(self):
+    def print(self, compute=False):
         """Prints the statistics of the Cutflow"""
 
-        if self._delayed_mode:
+        if self._delayed_mode and not compute:
+            warnings.warn(
+                "This will compute dask_awkward arrays. If you really want to do this now, call print(compute=True)"
+            )
+            return
+
+        if self._delayed_mode and compute:
             self._nevonecut, self._nevcutflow = dask.compute(
                 self._nevonecut, self._nevcutflow
             )