Refactor

Author: nalepae

pandarallel/data_types/dataframe.py

@@ -2,9 +2,8 @@
 
 import pandas as pd
 
-from ..utils import chunk
+from ..utils import chunk, get_axis_int
 from .generic import DataType
-from ..utils import _get_axis_int, _opposite_axis_int
 
 
 class DataFrame:
@@ -15,8 +14,8 @@ def get_chunks(
         ) -> Iterator[pd.DataFrame]:
             user_defined_function_kwargs = kwargs["user_defined_function_kwargs"]
 
-            axis_int = _get_axis_int(user_defined_function_kwargs)
-            opposite_axis_int = _opposite_axis_int(axis_int)
+            axis_int = get_axis_int(user_defined_function_kwargs)
+            opposite_axis_int = 1 - axis_int
 
             for chunk_ in chunk(data.shape[opposite_axis_int], nb_workers):
                 yield data.iloc[chunk_] if axis_int == 1 else data.iloc[:, chunk_]
@@ -36,17 +35,16 @@ def work(
             )
 
         @staticmethod
-        def get_reduce_extra(data: Any, user_defined_function_kwargs) -> Dict[str, Any]:
-            return {"axis": _get_axis_int(user_defined_function_kwargs)}
+        def get_reduce_extra(
+            data: Any, user_defined_function_kwargs: Dict[str, Any]
+        ) -> Dict[str, Any]:
+            return {"axis": get_axis_int(user_defined_function_kwargs)}
 
         @staticmethod
         def reduce(
             datas: Iterable[pd.DataFrame], extra: Dict[str, Any]
         ) -> pd.DataFrame:
-            if isinstance(datas[0], pd.Series):
-                axis = 0
-            else:
-                axis = _opposite_axis_int(extra["axis"])
+            axis = 0 if isinstance(datas[0], pd.Series) else 1 - extra["axis"]
             return pd.concat(datas, copy=False, axis=axis)
 
     class ApplyMap(DataType):

pandarallel/data_types/dataframe_groupby.py

@@ -40,7 +40,9 @@ def compute_result(
             return [compute_result(key, df) for key, df in data]
 
         @staticmethod
-        def get_reduce_extra(data: PandasDataFrameGroupBy, user_defined_function_kwargs) -> Dict[str, Any]:
+        def get_reduce_extra(
+            data: PandasDataFrameGroupBy, user_defined_function_kwargs: Dict[str, Any]
+        ) -> Dict[str, Any]:
             return {"df_groupby": data}
 
         @staticmethod

pandarallel/data_types/generic.py

@@ -24,7 +24,9 @@ def work(
         ...
 
     @staticmethod
-    def get_reduce_extra(data: Any, user_defined_function_kwargs) -> Dict[str, Any]:
+    def get_reduce_extra(
+        data: Any, user_defined_function_kwargs: Dict[str, Any]
+    ) -> Dict[str, Any]:
         return dict()
 
     @staticmethod

pandarallel/utils.py

@@ -1,6 +1,6 @@
 import itertools
 from enum import Enum
-from typing import List, Tuple
+from typing import Any, Dict, List, Tuple
 
 import pandas as pd
 from pandas import DataFrame, Index
@@ -87,17 +87,14 @@ def get_pandas_version() -> Tuple[int, int]:
     return int(major_str), int(minor_str)
 
 
-def _get_axis_int(user_defined_function_kwargs):
+def get_axis_int(user_defined_function_kwargs: Dict[str, Any]):
     axis = user_defined_function_kwargs.get("axis", 0)
 
     if axis not in {0, 1, "index", "columns"}:
         raise ValueError(f"No axis named {axis} for object type DataFrame")
 
-    axis_int = {0: 0, 1: 1, "index": 0, "columns": 1}[axis]
-    return axis_int
+    return {0: 0, 1: 1, "index": 0, "columns": 1}[axis]
 
-def _opposite_axis_int(axis: int):
-    return 1 - axis
 
 class WorkerStatus(int, Enum):
     Running = 0