Minor final optimizations.

yohplala · Sep 4, 2024 · 3e1bbaf · 3e1bbaf
1 parent 37b4f39
commit 3e1bbaf
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 10 deletions.
diff --git a/oups/aggstream/aggstream.py b/oups/aggstream/aggstream.py
@@ -730,16 +730,18 @@ def _post_n_write_agg_chunks(
                         f"key '{key[1]}' for snapshots but 'post()' function "
                         "only returning one result.",
                     )
-                # Set to None 'bin_res' and 'snap_res'.
-                # This allows to catch possible mistake in 'key' parameter.
-                snap_res = None
-            bin_res = None
+                # Set to None 'bin_res' and 'snap_res' to catch possible
+                # mistake in 'key' parameter (finally commented out).
+                # snap_res = None
+            # bin_res = None
         elif snap_res is None or isinstance(key, tuple):
             # Case only 'bin_res' is recorded or both 'bin_res' and 'snap_res'.
-            main_res, bin_res = bin_res, None
+            # main_res, bin_res = bin_res, None
+            main_res = bin_res
         else:
             # Case only 'snap_res' is recorded, and not 'bin_res'.
-            main_res, bin_res, snap_res = snap_res, None, None
+            # main_res, bin_res, snap_res = snap_res, None, None
+            main_res = snap_res
     else:
         not_null_res = False
         main_res = None

diff --git a/tests/test_aggstream/test_aggstream_advanced.py b/tests/test_aggstream/test_aggstream_advanced.py
@@ -1269,7 +1269,3 @@ def post(buffer: dict, bin_res: DataFrame, snap_res: DataFrame):
     seed_df = DataFrame({ordered_on: ts, val: rand_ints})
     with pytest.raises(ValueError, match="^not possible to have key 'agg_10T_bin'"):
         as_.agg(seed=seed_df, trim_start=False, discard_last=False, final_write=True)
-
-
-# TODO: Check error message when there is a post with both bin_key and snap_key, but only one dataframe is returned
-# TODO: in the end, clean 'aggstream': not set bin_res or snap_res to None if not necessary.