Fix Ensemble.batch(on=None/str)

Closes #202
lincc-frameworks · Aug 24, 2023 · c6f367c · c6f367c
1 parent b675119
commit c6f367c
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/tape/ensemble.py b/src/tape/ensemble.py
@@ -728,6 +728,8 @@ def batch(self, func, *args, meta=None, use_map=True, compute=True, on=None, **k
 
         if on is None:
             on = self._id_col  # Default grouping is by id_col
+        if isinstance(on, str):
+            on = [on]  # Convert to list if only one column is passed
 
         # Handle object columns to group on
         source_cols = list(self._source.columns)