created separate config file for managing parameters of the BankchurnTransformer

Author: wegar-2

moddata/pipeline/bankchurn_pipeline.py

@@ -4,6 +4,7 @@
 
 from moddata.extractor.bankchurn_extractor import BankchurnExtractor
 from moddata.transformer.bankchurn_transformer import BankchurnTransformer
+from moddata.src.constants import EncodingAndScalingModelType
 
 
 class BankchurnPipeline:
@@ -12,12 +13,15 @@ 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
         self._transformer: Final[BankchurnTransformer] = (
             BankchurnTransformer(
                 train_size=train_size,
-                random_state=random_state
+                random_state=random_state,
+                encoding_and_scaling_model_type=self._encoding_and_scaling_model_type
             )
         )
 

moddata/src/config.py

@@ -0,0 +1,25 @@
+"""Stores pydantic-style configuration classes.
+
+These are mainly bundles of values used not to clutter __init__-s
+"""
+
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from moddata.src.constants import EncodingAndScalingModelType
+
+
+class BankchurnTransformerConfig(BaseModel):
+    """
+    train_size: parameter passed to the train_test_split method
+    used to create train and test datasets
+    random_state: analogous to train_size
+    encoding_and_scaling_model_type: Literal, defines what type of
+    model data should be prepared for
+    """
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    train_size: float | int
+    random_state: Optional[int] = None,
+    encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = None

moddata/transformer/bankchurn_transformer.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Final
 
 import pandas as pd
 from sklearn.model_selection import train_test_split
@@ -9,29 +9,13 @@
 
 from moddata.src.constants import EncodingAndScalingModelType
 from moddata.sklearn_extensions.log_standard_scaler import LogStandardScaler
+from moddata.src.config import BankchurnTransformerConfig
 
 
 class BankchurnTransformer:
 
-    def __init__(
-            self,
-            train_size: float | int,
-            random_state: Optional[int] = None,
-            encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = None,
-    ):
-        """
-
-        Args:
-            train_size: parameter passed to the train_test_split method
-            used to create train and test datasets
-            random_state: analogous to train_size
-            encoding_and_scaling_model_type: Literal, defines what type of
-            model data should be prepared for
-        """
-        self._train_size: float | int = train_size
-        self._random_state: Optional[int] = random_state
-        self._encoding_and_scaling_model_type: Optional[EncodingAndScalingModelType] = (
-            encoding_and_scaling_model_type)
+    def __init__(self, config: BankchurnTransformerConfig):
+        self._config: Final[BankchurnTransformerConfig] = config
 
     @staticmethod
     def _ohe_gender_encoder() -> OneHotEncoder:
@@ -101,14 +85,13 @@ def transform(
     ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
         X, y = data
         X_train, X_test, y_train, y_test = train_test_split(
-            X,
-            y,
-            train_size=self._train_size,
-            random_state=self._random_state
+            X, y,
+            train_size=self._config.train_size,
+            random_state=self._config.random_state
         )
-        if self._encoding_and_scaling_model_type is not None:
+        if self._config.encoding_and_scaling_model_type is not None:
             col_trfm: ColumnTransformer = self._get_column_transformer(
-                encoding_and_scaling_model_type=self._encoding_and_scaling_model_type
+                encoding_and_scaling_model_type=self._config.encoding_and_scaling_model_type
             )
             X_train, y_train = col_trfm.fit_transform(X=X_train, y=y_train)
             X_test, y_test = col_trfm.transform(X=X_test, y=y_test)