style: address formatting, linting and type checking issues

Author: stevenhsd

src/dve/core_engine/backends/base/backend.py

@@ -3,7 +3,7 @@
 import logging
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Mapping, MutableMapping
+from collections.abc import MutableMapping
 from typing import Any, ClassVar, Generic, Optional
 
 from pyspark.sql import DataFrame, SparkSession
@@ -67,6 +67,7 @@ def load_reference_data(
         reference_entity_config: dict[EntityName, ReferenceConfigUnion],
         submission_info: Optional[SubmissionInfo],
     ) -> BaseRefDataLoader[EntityType]:
+        """Supply configured reference data loader for use with business rules"""
         raise NotImplementedError()
 
     @abstractmethod

src/dve/core_engine/backends/base/reference_data.py

@@ -208,8 +208,8 @@ def __getitem__(self, key: EntityName) -> EntityType:
             try:
                 config = self.reference_entity_config[key]
                 return self.load_entity(entity_name=key, config=config)
-            except TypeError:
-                raise NoRefDataConfigSupplied()
+            except TypeError as err:
+                raise NoRefDataConfigSupplied() from err
             except Exception as err:
                 raise MissingRefDataEntity(entity_name=key) from err
 

src/dve/core_engine/backends/exceptions.py

@@ -118,9 +118,10 @@ def get_message_preamble(self) -> str:
         """
         return f"Missing reference data entity {self.entity_name!r}"
 
+
 class NoRefDataConfigSupplied(BackendError):
     """An error raised when trying to load a refdata entity when no refdata
-       config has been supplied.
+    config has been supplied.
 
     """
 
@@ -129,7 +130,7 @@ def __init__(self, *args: object) -> None:
 
     def get_message_preamble(self) -> EntityName:
         """Message for logging purposes"""
-        return f"Refdata loader not supplied with refdata config - unable to load refdata entities"
+        return "Refdata loader not supplied with refdata config - unable to load refdata entities"
 
 
 class ConstraintError(ValueError, BackendErrorMixin):

src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py

@@ -411,7 +411,7 @@ def get_duckdb_cast_statement_from_annotation(
             stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{date_regex}') THEN TRY_CAST(TRIM({quoted_name}) as DATE) ELSE NULL END"  # pylint: disable=C0301
             return stmt
         if issubclass(type_, time):
-            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{time_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIME) ELSE NULL END" # pylint: disable=C0301
+            stmt = rf"CASE WHEN REGEXP_MATCHES(TRIM({quoted_name}), '{time_regex}') THEN TRY_CAST(TRIM({quoted_name}) as TIME) ELSE NULL END"  # pylint: disable=C0301
             return stmt
         duck_type = get_duckdb_type_from_annotation(type_)
         if duck_type:

src/dve/core_engine/backends/implementations/duckdb/reference_data.py

@@ -1,14 +1,11 @@
 """A reference data loader for duckdb."""
 
-from typing import Optional
-
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pyarrow import ipc  # type: ignore
 
 from dve.core_engine.backends.base.reference_data import (
     BaseRefDataLoader,
     ReferenceConfig,
-    ReferenceConfigUnion,
     ReferenceTable,
     mark_refdata_file_extension,
 )
@@ -19,18 +16,18 @@
 # pylint: disable=too-few-public-methods
 class DuckDBRefDataLoader(BaseRefDataLoader[DuckDBPyRelation]):
     """A reference data loader using already existing DuckDB tables.
-       reference_entity_config and dataset_config_uri (if config uses relative paths)
-       should be supplied using setter methods for the dataset being processed before running."""
+    reference_entity_config and dataset_config_uri (if config uses relative paths)
+    should be supplied using setter methods for the dataset being processed before running."""
 
     def __init__(
         self,
         connection: DuckDBPyConnection,
         reference_data_config: dict[EntityName, ReferenceConfig],
         dataset_config_uri: URI,
-        **kwargs
+        **kwargs,
     ) -> None:
-        super().__init__(reference_data_config, dataset_config_uri,**kwargs)
-        
+        super().__init__(reference_data_config, dataset_config_uri, **kwargs)
+
         self.connection = connection
 
         if not self.connection:

src/dve/core_engine/backends/implementations/spark/backend.py

@@ -17,7 +17,7 @@
 from dve.core_engine.loggers import get_child_logger, get_logger
 from dve.core_engine.models import SubmissionInfo
 from dve.core_engine.type_hints import URI, EntityName, EntityParquetLocations
-from dve.parser.file_handling import get_resource_exists, joinuri, get_parent
+from dve.parser.file_handling import get_resource_exists, joinuri
 
 
 class SparkBackend(BaseBackend[DataFrame]):
@@ -50,18 +50,22 @@ def __init__(
                 logger=get_child_logger("SparkStepImplementations", logger)
             )
         super().__init__(contract, steps, logger, **kwargs)
-    
-    def load_reference_data(self,
+
+    def load_reference_data(
+        self,
         reference_entity_config: dict[EntityName, ReferenceConfigUnion],
-        submission_info: Optional[SubmissionInfo],):
+        submission_info: Optional[SubmissionInfo],
+    ):
         """Load the reference data as specified in the reference entity config."""
-        sub_info_entity: Optional[EntityType] = None
+        sub_info_entity: Optional[DataFrame] = None
         if submission_info:
             sub_info_entity = self.convert_submission_info(submission_info)
 
-        reference_data_loader = SparkRefDataLoader(spark=self.spark_session,
-                                                   reference_data_config=reference_entity_config,
-                                                   dataset_config_uri=self.dataset_config_uri)
+        reference_data_loader = SparkRefDataLoader(
+            spark=self.spark_session,
+            reference_data_config=reference_entity_config,
+            dataset_config_uri=self.dataset_config_uri, # type: ignore
+        )
         if sub_info_entity is not None:
             reference_data_loader.entity_cache["dve_submission_info"] = sub_info_entity
 

src/dve/core_engine/backends/implementations/spark/reference_data.py

@@ -1,8 +1,6 @@
 # pylint: disable=no-member
 """A reference data loader for Spark."""
 
-from typing import Optional
-
 from pyspark.sql import DataFrame, SparkSession
 
 from dve.core_engine.backends.base.reference_data import (
@@ -18,8 +16,8 @@
 # pylint: disable=too-few-public-methods
 class SparkRefDataLoader(BaseRefDataLoader[DataFrame]):
     """A reference data loader using already existing Apache Spark Tables.
-       reference_entity_config and dataset_config_uri (if config uses relative paths)
-       should be supplied using setter methods for the dataset being processed before running."""
+    reference_entity_config and dataset_config_uri (if config uses relative paths)
+    should be supplied using setter methods for the dataset being processed before running."""
 
     def __init__(
         self,

src/dve/pipeline/duckdb_pipeline.py

@@ -5,6 +5,7 @@
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
+import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
 from dve.core_engine.backends.implementations.duckdb.auditing import DDBAuditingManager
 from dve.core_engine.backends.implementations.duckdb.contract import DuckDBDataContract
@@ -14,7 +15,6 @@
 from dve.core_engine.models import SubmissionInfo
 from dve.core_engine.type_hints import URI
 from dve.pipeline.pipeline import BaseDVEPipeline
-import dve.parser.file_handling as fh
 
 
 # pylint: disable=abstract-method
@@ -47,13 +47,16 @@ def __init__(
             logger,
         )
 
-    def get_reference_data_loader(self,
-                           reference_data_config: dict[str, ReferenceConfig],
-                           **kwargs) -> BaseRefDataLoader[DuckDBPyRelation]:
-        return DuckDBRefDataLoader(connection=self._connection,
-                                   reference_data_config=reference_data_config,
-                                   dataset_config_uri=fh.get_parent(self._rules_path),
-                                   **kwargs)
+    def get_reference_data_loader(
+        self, reference_data_config: dict[str, ReferenceConfig], **kwargs
+    ) -> DuckDBRefDataLoader:
+        return DuckDBRefDataLoader(
+            connection=self._connection,
+            reference_data_config=reference_data_config,
+            dataset_config_uri=fh.get_parent(self._rules_path), # type: ignore
+            **kwargs
+        )
+
     # pylint: disable=arguments-differ
     def write_file_to_parquet(  # type: ignore
         self, submission_file_uri: URI, submission_info: SubmissionInfo, output: URI

src/dve/pipeline/pipeline.py

@@ -112,13 +112,12 @@ def step_implementations(self) -> Optional[BaseStepImplementations[EntityType]]:
     def get_entity_count(entity: EntityType) -> int:
         """Get a row count of an entity stored as parquet"""
         raise NotImplementedError()
-    
-    def get_reference_data_loader(self,
-                           reference_data_config: dict[EntityName, ReferenceConfig],
-                           **kwargs) -> BaseRefDataLoader[EntityType]:
+
+    def get_reference_data_loader(
+        self, reference_data_config: dict[EntityName, ReferenceConfig], **kwargs
+    ) -> BaseRefDataLoader:
         """Get reference data loader if required for business rules"""
         raise NotImplementedError()
-        
 
     def get_submission_status(
         self, step_name: DVEStageName, submission_id: str
@@ -533,7 +532,7 @@ def data_contract_step(
 
         return processed_files, failed_processing
 
-    def apply_business_rules(   # pylint: disable=R0914
+    def apply_business_rules(  # pylint: disable=R0914
         self, submission_info: SubmissionInfo, submission_status: Optional[SubmissionStatus] = None
     ) -> tuple[SubmissionInfo, SubmissionStatus]:
         """Apply the business rules to a given submission, the submission may have failed at the
@@ -559,7 +558,7 @@ def apply_business_rules(   # pylint: disable=R0914
             self._processed_files_path, submission_info.submission_id
         )
         ref_data = config.get_reference_data_config()
-        reference_data = self.get_reference_data_loader(reference_data_config=ref_data)
+        reference_data: BaseRefDataLoader = self.get_reference_data_loader(reference_data_config=ref_data)
         rules = config.get_rule_metadata()
         entities = {}
         contract = fh.joinuri(
@@ -585,10 +584,7 @@ def apply_business_rules(   # pylint: disable=R0914
         key_fields = {model: conf.reporting_fields for model, conf in model_config.items()}
 
         _errors_uri, rules_success = self.step_implementations.apply_rules(  # type: ignore
-            working_directory,
-            entity_manager,
-            rules,
-            key_fields
+            working_directory, entity_manager, rules, key_fields
         )
 
         rule_messages = load_feedback_messages(

src/dve/pipeline/spark_pipeline.py

@@ -6,6 +6,7 @@
 
 from pyspark.sql import DataFrame, SparkSession
 
+import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
 from dve.core_engine.backends.implementations.spark.auditing import SparkAuditingManager
 from dve.core_engine.backends.implementations.spark.contract import SparkDataContract
@@ -16,7 +17,6 @@
 from dve.core_engine.type_hints import URI
 from dve.pipeline.pipeline import BaseDVEPipeline
 from dve.pipeline.utils import SubmissionStatus, unpersist_all_rdds
-import dve.parser.file_handling as fh
 
 
 # pylint: disable=abstract-method
@@ -48,14 +48,16 @@ def __init__(
             job_run_id,
             logger,
         )
-    
-    def get_reference_data_loader(self,
-                           reference_data_config: dict[str, ReferenceConfig],
-                           **kwargs) -> BaseRefDataLoader[DataFrame]:
-        return SparkRefDataLoader(spark=self._spark,
-                                  reference_data_config=reference_data_config,
-                                  dataset_config_uri=fh.get_parent(self._rules_path),
-                                  **kwargs)
+
+    def get_reference_data_loader(
+        self, reference_data_config: dict[str, ReferenceConfig], **kwargs
+    ) -> SparkRefDataLoader:
+        return SparkRefDataLoader(
+            spark=self._spark,
+            reference_data_config=reference_data_config,
+            dataset_config_uri=fh.get_parent(self._rules_path), # type: ignore
+            **kwargs
+        )
 
     # pylint: disable=arguments-differ
     def write_file_to_parquet(  # type: ignore