fix lint errors

Author: atulikumwenayo

src/datacustomcode/common_config.py

@@ -1,12 +1,26 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import os
-import yaml
+from typing import Any
+
 from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
 )
-
-from typing import Any
+import yaml
 
 DEFAULT_CONFIG_NAME = "config.yaml"
 
@@ -18,7 +32,8 @@ def default_config_file() -> str:
 class ForceableConfig(BaseModel):
     force: bool = Field(
         default=False,
-        description="If True, this takes precedence over parameters passed to the initializer of the client",
+        description="If True, this takes precedence over parameters passed to the "
+        "initializer of the client",
     )
 
 
@@ -34,6 +49,9 @@ class BaseObjectConfig(ForceableConfig):
 
 
 class BaseConfig(BaseModel):
+    def update(self, other: Any) -> "BaseConfig":
+        raise NotImplementedError("Subclasses must implement update method")
+
     def load(self, config_path: str) -> "BaseConfig":
         """Load configuration from a YAML file and merge with existing config"""
         with open(config_path, "r") as f:

src/datacustomcode/config.py

@@ -25,8 +25,13 @@
     cast,
 )
 
-from pydantic import (
-    Field,
+from pydantic import Field
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    ForceableConfig,
+    default_config_file,
 )
 
 # This lets all readers and writers to be findable via config
@@ -37,8 +42,6 @@
 from datacustomcode.proxy.base import BaseProxyAccessLayer
 from datacustomcode.proxy.client.base import BaseProxyClient  # noqa: TCH002
 from datacustomcode.spark.base import BaseSparkSessionProvider
-from datacustomcode.common_config import ForceableConfig, BaseObjectConfig, BaseConfig, default_config_file
-
 
 if TYPE_CHECKING:
     from pyspark.sql import SparkSession
@@ -49,6 +52,7 @@
 
 class AccessLayerObjectConfig(BaseObjectConfig, Generic[_T]):
     type_base: ClassVar[Type[BaseDataAccessLayer]] = BaseDataAccessLayer
+
     def to_object(self, spark: SparkSession) -> _T:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
         return cast(_T, type_(spark=spark, **self.options))
@@ -75,14 +79,17 @@ class SparkConfig(ForceableConfig):
 
 class ProxyAccessLayerObjectConfig(BaseObjectConfig, Generic[_PX]):
     """Config for proxy clients that take no constructor args (e.g. no spark)."""
+
     type_base: ClassVar[Type[BaseProxyAccessLayer]] = BaseProxyAccessLayer
+
     def to_object(self) -> _PX:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
         return cast(_PX, type_(**self.options))
 
 
 class SparkProviderConfig(BaseObjectConfig, Generic[_P]):
     type_base: ClassVar[Type[BaseSparkSessionProvider]] = BaseSparkSessionProvider
+
     def to_object(self) -> _P:
         type_ = self.type_base.subclass_from_config_name(self.type_config_name)
         return cast(_P, type_(**self.options))
@@ -126,6 +133,7 @@ def merge(
         )
         return self
 
+
 """Global config object.
 
 This is the object that makes config accessible globally and globally mutable.

src/datacustomcode/einstein_predictions/__init__.py

@@ -19,4 +19,4 @@
 __all__ = [
     "EinsteinPredictions",
     "DefaultEinsteinPredictions",
-]
+]

src/datacustomcode/einstein_predictions/base.py

@@ -15,9 +15,13 @@
 
 from abc import ABC, abstractmethod
 
-from datacustomcode.einstein_predictions.types import (PredictionRequest, PredictionResponse)
+from datacustomcode.einstein_predictions.types import (
+    PredictionRequest,
+    PredictionResponse,
+)
 from datacustomcode.mixin import UserExtendableNamedConfigMixin
 
+
 class EinsteinPredictions(ABC, UserExtendableNamedConfigMixin):
     CONFIG_NAME: str
 

src/datacustomcode/einstein_predictions/impl/default.py

@@ -1,4 +1,3 @@
-
 # Copyright (c) 2025, Salesforce, Inc.
 # SPDX-License-Identifier: Apache-2
 #
@@ -17,9 +16,10 @@
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
 from datacustomcode.einstein_predictions.types import (
     PredictionRequest,
-    PredictionResponse
+    PredictionResponse,
 )
 
+
 class DefaultEinsteinPredictions(EinsteinPredictions):
     CONFIG_NAME = "DefaultEinsteinPredictions"
 
@@ -28,8 +28,8 @@ def __init__(self, **kwargs):
 
     def predict(self, request: PredictionRequest) -> PredictionResponse:
         return PredictionResponse(
-              version="v1",
-              prediction_type=request.prediction_type,
-              status_code=200,
-              data={"results": [{"prediction": {"predictedValue": "1"}}]}
-        )
+            version="v1",
+            prediction_type=request.prediction_type,
+            status_code=200,
+            data={"results": [{"prediction": {"predictedValue": "1"}}]},
+        )

src/datacustomcode/einstein_predictions/types.py

@@ -14,18 +14,19 @@
 # limitations under the License.
 
 from enum import Enum, unique
+from typing import (
+    Any,
+    Dict,
+    Literal,
+    Optional,
+)
+
 from pydantic import (
     BaseModel,
     Field,
     model_validator,
 )
 
-from typing import (
-    Literal,
-    Optional,
-    Dict,
-    Any
-)
 
 @unique
 class PredictionType(Enum):
@@ -35,37 +36,51 @@ class PredictionType(Enum):
     MULTI_OUTCOME = 4
     BINARY_CLASSIFICATION = 5
 
+
 class PredictionColumn(BaseModel):
     column_name: str = Field(min_length=1, description="Column name")
-    string_values: Optional[list[str]] = Field(default = None, min_length=1, description="Column string values")
-    double_values: Optional[list[float]] = Field(default = None, min_length=1, description="Column double values")
-    boolean_values: Optional[list[bool]] = Field(default = None, min_length=1, description="Column boolean values")
-    date_values: Optional[list[str]] = Field(default = None, min_length=1, description="Column date values")
-    datetime_values: Optional[list[str]] = Field(default = None, min_length=1, description="Column datetime values")
+    string_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column string values"
+    )
+    double_values: Optional[list[float]] = Field(
+        default=None, min_length=1, description="Column double values"
+    )
+    boolean_values: Optional[list[bool]] = Field(
+        default=None, min_length=1, description="Column boolean values"
+    )
+    date_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column date values"
+    )
+    datetime_values: Optional[list[str]] = Field(
+        default=None, min_length=1, description="Column datetime values"
+    )
 
-    @model_validator(mode='after')
+    @model_validator(mode="after")
     def validate_exactly_one_value_type(self):
-        set_count = sum([
-            self.string_values is not None,
-            self.double_values is not None,
-            self.boolean_values is not None,
-            self.date_values is not None,
-            self.datetime_values is not None
-        ])
+        set_count = sum(
+            [
+                self.string_values is not None,
+                self.double_values is not None,
+                self.boolean_values is not None,
+                self.date_values is not None,
+                self.datetime_values is not None,
+            ]
+        )
 
         if set_count != 1:
             raise ValueError("Exactly one value type must be set")
 
         return self
 
+
 class PredictionColumBuilder:
     def __init__(self) -> None:
-        self._column_name: str = None
-        self._string_values: list[str] = None
-        self._double_values: list[float] = None
-        self._boolean_values: list[bool] = None
-        self._date_values: list[str] = None
-        self._datetime_values: list[str] = None
+        self._column_name: Optional[str] = None
+        self._string_values: Optional[list[str]] = None
+        self._double_values: Optional[list[float]] = None
+        self._boolean_values: Optional[list[bool]] = None
+        self._date_values: Optional[list[str]] = None
+        self._datetime_values: Optional[list[str]] = None
 
     def set_column_name(self, column_name: str) -> "PredictionColumBuilder":
         self._column_name = column_name
@@ -79,45 +94,59 @@ def set_double_values(self, double_values: list[float]) -> "PredictionColumBuild
         self._double_values = double_values
         return self
 
-    def set_boolean_values(self, boolean_values: list[bool]) -> "PredictionColumBuilder":
+    def set_boolean_values(
+        self, boolean_values: list[bool]
+    ) -> "PredictionColumBuilder":
         self._boolean_values = boolean_values
         return self
 
     def set_date_values(self, date_values: list[str]) -> "PredictionColumBuilder":
         self._date_values = date_values
         return self
 
-    def set_datetime_values(self, datetime_values: list[str]) -> "PredictionColumBuilder":
+    def set_datetime_values(
+        self, datetime_values: list[str]
+    ) -> "PredictionColumBuilder":
         self._datetime_values = datetime_values
         return self
 
     def build(self) -> PredictionColumn:
         return PredictionColumn(
-            column_name = self._column_name,
-            string_values = self._string_values,
-            double_values = self._double_values,
-            boolean_values = self._boolean_values,
-            date_values = self._date_values,
-            datetime_values = self._datetime_values
+            column_name=self._column_name,
+            string_values=self._string_values,
+            double_values=self._double_values,
+            boolean_values=self._boolean_values,
+            date_values=self._date_values,
+            datetime_values=self._datetime_values,
         )
 
+
 class PredictionRequest(BaseModel):
     version: Literal["v1"] = Field(
         default="v1", description="API version, must be 'v1'"
     )
     prediction_type: PredictionType = Field(description="Prediction type")
-    model_api_name: str = Field(min_length=1, description="API name of the model to use")
-    prediction_columns: list[PredictionColumn] = Field(min_length=1, description="List of prediction columns")
-    settings: Optional[Dict[str, Any]] = Field(default=None, description="Settings for the prediction request")
+    model_api_name: str = Field(
+        min_length=1, description="API name of the model to use"
+    )
+    prediction_columns: list[PredictionColumn] = Field(
+        min_length=1, description="List of prediction columns"
+    )
+    settings: Optional[Dict[str, Any]] = Field(
+        default=None, description="Settings for the prediction request"
+    )
+
 
 class PredictionRequestBuilder:
     def __init__(self) -> None:
-        self._prediction_type: PredictionType = None
-        self._model_api_name: str = None
+        self._prediction_type: Optional[PredictionType] = None
+        self._model_api_name: Optional[str] = None
         self._prediction_columns: list[PredictionColumn] = []
-        self._settings: Dict[str, Any] = None
+        self._settings: Optional[Dict[str, Any]] = None
 
-    def set_prediction_type(self, prediction_type: PredictionType) -> "PredictionRequestBuilder":
+    def set_prediction_type(
+        self, prediction_type: PredictionType
+    ) -> "PredictionRequestBuilder":
         self._prediction_type = prediction_type
         return self
 
@@ -126,8 +155,7 @@ def set_model(self, model_api_name: str) -> "PredictionRequestBuilder":
         return self
 
     def set_prediction_columns(
-        self,
-        prediction_columns: list[PredictionColumn]
+        self, prediction_columns: list[PredictionColumn]
     ) -> "PredictionRequestBuilder":
         self._prediction_columns = prediction_columns
         return self
@@ -141,9 +169,10 @@ def build(self) -> PredictionRequest:
             prediction_type=self._prediction_type,
             model_api_name=self._model_api_name,
             prediction_columns=self._prediction_columns,
-            settings=self._settings
+            settings=self._settings,
         )
-    
+
+
 class PredictionResponse(BaseModel):
     version: Literal["v1"] = Field(default="v1", description="API version")
     prediction_type: PredictionType = Field(description="Prediction type")
@@ -153,4 +182,3 @@ class PredictionResponse(BaseModel):
     @property
     def is_success(self) -> bool:
         return self.status_code == 200
-