[read_csv] -update chunk function arg logic

Author: Alex0x4b

src/akutils/pandas_read_files.py

@@ -12,8 +12,7 @@
 )
 from akutils.utils_functions import (
     timeit,
-    sanitize_function_args_from_locals,
-    control_if_usecols_exist_in_df
+    contruct_function_args_from_locals,
 )
 from akutils.os import list_files_from_dir, warn
 
@@ -22,6 +21,7 @@
 def read_csv_in_chunks(
     filepath_or_buffer: FilePath | ReadCsvBuffer[bytes] | ReadCsvBuffer[str],
     chunk_func: Callable | None = None,
+    chunk_func_kwarg=None,
     chunksize: int = 10**6,
     dtype: DtypeArg | None = "string",
     **kwargs
@@ -62,17 +62,17 @@ def filter_on_countries(df, countries):
             chunksize=5
         )
     """
+    if chunk_func_kwarg is None:
+        chunk_func_kwarg = {}
     print(f"File: {filepath_or_buffer}")
     locals_args = locals()  # get all args passed in the function
-    read_csv_args = sanitize_function_args_from_locals(pd.read_csv, locals_args)
-    read_csv_args = control_if_usecols_exist_in_df(**read_csv_args)
+    read_csv_args = contruct_function_args_from_locals(pd.read_csv, locals_args)
 
     df = pd.DataFrame()
     counter = 0
     for chunk in pd.read_csv(**read_csv_args):
         print(f"Chunk number => {counter}")
-        chunk_func_kwarg = sanitize_function_args_from_locals(chunk_func, locals_args)
-        chunk = chunk_func(chunk, **chunk_func_kwarg) if chunk_func else chunk
+        chunk = chunk_func(df=chunk, **chunk_func_kwarg) if chunk_func else chunk
         df = pd.concat([df, chunk], axis=0, ignore_index=True)
         counter += 1
     return df

src/akutils/tests/unit/test_pandas_read_files.py

@@ -28,15 +28,15 @@ def test_read_csv_in_chunk_with_filtered_function(self):
             columns=["col1", "col2", "country"]
         )
 
-        def filter_chunk(df_chunk: pd.DataFrame, countries: list) -> pd.DataFrame:
-            return df_chunk[df_chunk["country"].isin(countries)]
+        def filter_chunk(df: pd.DataFrame, countries: list) -> pd.DataFrame:
+            return df[df["country"].isin(countries)]
 
         file_path = PATH_TO_AKUTILS_PKG / "tests" / "_fixtures" / "sales.csv"
         df = ak.read_csv_in_chunks(
             file_path,
             chunk_func=filter_chunk,
+            chunk_func_kwarg={"countries": ["Spain"]},
             chunksize=5,
-            countries=["Spain"],
             sep=";",
             dtype=None
         )
@@ -51,21 +51,23 @@ def test_read_csv_in_chunk_with_add_cols_function(self):
         df_expected["new_col"] = (df_expected["col1"] + df_expected["col2"]) * 2.4 + 151
 
         def add_new_col_to_chunk(
-            df_chunk: pd.DataFrame,
+            df: pd.DataFrame,
             factor: float,
             constant: int
         ) -> pd.DataFrame:
-            df_chunk["new_col"] = (
-                (df_chunk["col1"] + df_chunk["col2"]) * factor + constant
+            df["new_col"] = (
+                (df["col1"] + df["col2"]) * factor + constant
             )
-            return df_chunk
+            return df
 
         file_path = PATH_TO_AKUTILS_PKG / "tests" / "_fixtures" / "sales.csv"
         df = ak.read_csv_in_chunks(
             file_path,
             chunk_func=add_new_col_to_chunk,
-            factor=2.4,
-            constant=151,
+            chunk_func_kwarg={
+                "factor": 2.4,
+                "constant": 151,
+            },
             chunksize=5,
             sep=";",
             dtype=None
@@ -134,8 +136,8 @@ def test_read_multiple_csv_from_dir_with_chunk_function(self):
         Case with chunk function (the dir containing also not text file extension)
         """
 
-        def filter_chunk(df_chunk: pd.DataFrame) -> pd.DataFrame:
-            return df_chunk[df_chunk["nb_sales"] <= 4]
+        def filter_chunk(df: pd.DataFrame) -> pd.DataFrame:
+            return df[df["nb_sales"] <= 4]
 
         df_expected = self.df_expected.copy()
         df_expected = df_expected[df_expected["nb_sales"] <= 4]

src/akutils/utils_functions.py

@@ -17,6 +17,22 @@ def timeit_wrapper(*args, **kwargs):
     return timeit_wrapper
 
 
+def contruct_function_args_from_locals(function, locals_args):
+    specified_args = {
+        key: value for key, value in locals_args.items() if key not in ["kwargs"]
+    }
+    additionnal_kwargs = locals_args["kwargs"]
+    all_args = dict(specified_args, **additionnal_kwargs)
+
+    # Filter on function alloewed args
+    function_args = {
+        key: value
+        for key, value in all_args.items()
+        if key in function.__code__.co_varnames
+    }
+    return function_args
+
+
 def sanitize_function_args_from_locals(function, locals_args):
     # Check if a function is passed, if not return empty dict
     if not hasattr(function, '__call__'):