diff --git a/oups/aggstream/aggstream.py b/oups/aggstream/aggstream.py index 5359815..496ff3a 100644 --- 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 index 77ed43b..8fd2b89 100644 --- 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.