Added config for the BankchurnPipeline and created aliases for tuples of DataFrames

Author: wegar-2

moddata/extractor/bankchurn_extractor.py

@@ -1,11 +1,11 @@
 import pandas as pd
 
 from moddata import load_data
-
+from moddata.src.constants import XyDataFrames
 
 class BankchurnExtractor:
 
-    def extract(self) -> tuple[pd.DataFrame, pd.DataFrame]:
+    def extract(self) -> XyDataFrames:
         data: pd.DataFrame = load_data(dataset="bankchurn")
         x: pd.DataFrame = data.loc[:, data.columns != "churn"]
         x = x.drop(columns=["customer_id"])

moddata/pipeline/bankchurn_pipeline.py

@@ -1,30 +1,19 @@
-from typing import Final, Optional
-
-import pandas as pd
+from typing import Final
 
 from moddata.extractor.bankchurn_extractor import BankchurnExtractor
 from moddata.transformer.bankchurn_transformer import BankchurnTransformer
-from moddata.src.constants import EncodingAndScalingModelType
+from moddata.src.constants import TrainTestXyDataFrames
+from moddata.src.config import BankchurnPipelineConfig
 
 
 class BankchurnPipeline:
 
-    def __init__(
-            self,
-            train_size: float | int,
-            random_state: Optional[int] = None,
-            encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = None
-    ):
-        self._random_state: Optional[int] = random_state
-        self._encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = encoding_and_scaling_model_type
+    def __init__(self, config: BankchurnPipelineConfig):
+        self._config: Final[BankchurnPipelineConfig] = config
         self._transformer: Final[BankchurnTransformer] = (
-            BankchurnTransformer(
-                train_size=train_size,
-                random_state=random_state,
-                encoding_and_scaling_model_type=self._encoding_and_scaling_model_type
-            )
+            BankchurnTransformer(config=self._config)
         )
 
-    def run(self) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    def run(self) -> TrainTestXyDataFrames:
         x, y = BankchurnExtractor().extract()
         return self._transformer.transform(data=(x, y))

moddata/src/config.py

@@ -3,12 +3,17 @@
 These are mainly bundles of values used not to clutter __init__-s
 """
 
-from typing import Optional
+from typing import Optional, TypeAlias
 
 from pydantic import BaseModel, ConfigDict
 
 from moddata.src.constants import EncodingAndScalingModelType
 
+__all__ = [
+    "BankchurnTransformerConfig",
+    "BankchurnPipelineConfig"
+]
+
 
 class BankchurnTransformerConfig(BaseModel):
     """
@@ -23,3 +28,6 @@ class BankchurnTransformerConfig(BaseModel):
     train_size: float | int
     random_state: Optional[int] = None,
     encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = None
+
+
+BankchurnPipelineConfig: TypeAlias = BankchurnTransformerConfig

moddata/src/constants.py

@@ -7,8 +7,15 @@
 
 from typing import TypeAlias, Literal
 
+import pandas as pd
 
 EncodingAndScalingModelType: TypeAlias = Literal[
     "tree_like",
     "other"
-]
+]
+
+XyDataFrames: TypeAlias = tuple[pd.DataFrame, pd.DataFrame]
+
+TrainTestXyDataFrames: TypeAlias = (
+    tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]
+)

moddata/transformer/bankchurn_transformer.py

@@ -10,6 +10,7 @@
 from moddata.src.constants import EncodingAndScalingModelType
 from moddata.sklearn_extensions.log_standard_scaler import LogStandardScaler
 from moddata.src.config import BankchurnTransformerConfig
+from moddata.src.constants import TrainTestXyDataFrames
 
 
 class BankchurnTransformer:
@@ -82,7 +83,7 @@ def _get_column_transformer(
     def transform(
             self,
             data: tuple[pd.DataFrame, pd.DataFrame]
-    ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    ) -> TrainTestXyDataFrames:
         X, y = data
         X_train, X_test, y_train, y_test = train_test_split(
             X, y,

tests/sklearn_extensions/test_log_standard_scaler.py

@@ -30,7 +30,6 @@ def test_log_standard_scaler_with_shift(make_log_normal_array):
 def test_log_standard_scaler_with_base(make_log_normal_array):
     lss: LogStandardScaler = LogStandardScaler(log_base=2)
     X_trfmd: np.ndarray = lss.fit_transform(X=make_log_normal_array)
-    # print(f"{X_trfmd=}")
     assert round(float(X_trfmd[0, 0]), 8) == -1.41835861
     assert X_trfmd.shape == (5, 1)
     assert round(float(X_trfmd[-1, 0]), 8) == -0.64893433