style: run black, isort and resolve linting issues

Author: stevenhsd

src/dve/common/error_utils.py

@@ -1,27 +1,29 @@
 """Utilities to support reporting"""
 
 import datetime as dt
-from itertools import chain
 import json
 import logging
+from collections.abc import Iterable
+from itertools import chain
 from multiprocessing import Queue
 from threading import Thread
-from typing import Iterable, Iterator, Optional, Union
+from typing import Optional, Union
 
-from dve.core_engine.message import UserMessage
-from dve.core_engine.loggers import get_logger
 import dve.parser.file_handling as fh
 from dve.core_engine.exceptions import CriticalProcessingError
+from dve.core_engine.loggers import get_logger
+from dve.core_engine.message import UserMessage
 from dve.core_engine.type_hints import URI, DVEStage, Messages
 
 
 def get_feedback_errors_uri(working_folder: URI, step_name: DVEStage) -> URI:
     """Determine the location of json lines file containing all errors generated in a step"""
     return fh.joinuri(working_folder, "errors", f"{step_name}_errors.jsonl")
 
+
 def get_processing_errors_uri(working_folder: URI) -> URI:
     """Determine the location of json lines file containing all processing
-       errors generated from DVE run"""
+    errors generated from DVE run"""
     return fh.joinuri(working_folder, "errors", "processing_errors.jsonl")
 
 
@@ -88,74 +90,85 @@ def dump_processing_errors(
         )
 
     with fh.open_stream(error_file, "a") as f:
-            f.write("\n".join([json.dumps(rec, default=str) for rec in processed]) + "\n")
-    
+        f.write("\n".join([json.dumps(rec, default=str) for rec in processed]) + "\n")
+
     return error_file
 
+
 def load_feedback_messages(feedback_messages_uri: URI) -> Iterable[UserMessage]:
+    """Load user messages from jsonl file"""
     if not fh.get_resource_exists(feedback_messages_uri):
         return
     with fh.open_stream(feedback_messages_uri) as errs:
         yield from (UserMessage(**json.loads(err)) for err in errs.readlines())
 
+
 def load_all_error_messages(error_directory_uri: URI) -> Iterable[UserMessage]:
-    return chain.from_iterable([load_feedback_messages(err_file) for err_file, _ in fh.iter_prefix(error_directory_uri) if err_file.endswith(".jsonl")])
+    "Load user messages from all jsonl files"
+    return chain.from_iterable(
+        [
+            load_feedback_messages(err_file)
+            for err_file, _ in fh.iter_prefix(error_directory_uri)
+            if err_file.endswith(".jsonl")
+        ]
+    )
+
 
 class BackgroundMessageWriter:
-    def __init__(self,
-                 working_directory: URI,
-                 dve_stage: DVEStage,
-                 key_fields: Optional[dict[str, list[str]]] = None,
-                 logger: Optional[logging.Logger] = None):
+    """Controls batch writes to error jsonl files"""
+
+    def __init__(
+        self,
+        working_directory: URI,
+        dve_stage: DVEStage,
+        key_fields: Optional[dict[str, list[str]]] = None,
+        logger: Optional[logging.Logger] = None,
+    ):
         self._working_directory = working_directory
         self._dve_stage = dve_stage
-        self._feedback_message_uri = get_feedback_errors_uri(self._working_directory, self._dve_stage)
+        self._feedback_message_uri = get_feedback_errors_uri(
+            self._working_directory, self._dve_stage
+        )
         self._key_fields = key_fields
         self.logger = logger or get_logger(type(self).__name__)
         self._write_thread = None
         self._queue = Queue()
-    
+
     @property
-    def write_queue(self):
+    def write_queue(self) -> Queue:  # type: ignore
+        """Queue for storing batches of messages to be written"""
         return self._queue
-    
+
     @property
-    def write_thread(self):
+    def write_thread(self) -> Thread:  # type: ignore
+        """Thread to write batches of messages to jsonl file"""
         if not self._write_thread:
             self._write_thread = Thread(target=self._write_process_wrapper)
         return self._write_thread
-        
-    
+
     def _write_process_wrapper(self):
         """Wrapper for dump feedback errors to run in background process"""
         while True:
             if msgs := self.write_queue.get():
-                dump_feedback_errors(self._working_directory, self._dve_stage, msgs, self._key_fields)
+                dump_feedback_errors(
+                    self._working_directory, self._dve_stage, msgs, self._key_fields
+                )
             else:
                 break
-    
+
     def __enter__(self) -> "BackgroundMessageWriter":
         self.write_thread.start()
         return self
-    
+
     def __exit__(self, exc_type, exc_value, traceback):
         if exc_type:
             self.logger.exception(
                 "Issue occured during background write process:",
-                exc_info=(exc_type, exc_value, traceback)
+                exc_info=(exc_type, exc_value, traceback),
             )
         self.write_queue.put(None)
         self.write_thread.join()
-        
-    
-
-def write_process_wrapper(working_directory: URI, *, queue: Queue, key_fields: Optional[dict[str, list[str]]] = None):
-    """Wrapper for dump feedback errors to run in background process"""
-    while True:
-        if msgs := queue.get():
-            dump_feedback_errors(fh.joinuri(working_directory, "data_contract"), msgs, key_fields)
-        else:
-            break
+
 
 def conditional_cast(value, primary_keys: list[str], value_separator: str) -> Union[list[str], str]:
     """Determines what to do with a value coming back from the error list"""

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

@@ -17,13 +17,7 @@
 from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful
 from dve.core_engine.loggers import get_logger
 from dve.core_engine.models import SubmissionInfo
-from dve.core_engine.type_hints import (
-    URI,
-    EntityLocations,
-    EntityName,
-    EntityParquetLocations,
-    Messages,
-)
+from dve.core_engine.type_hints import URI, EntityLocations, EntityName, EntityParquetLocations
 from dve.parser.file_handling.service import get_parent, joinuri
 
 
@@ -162,7 +156,9 @@ def apply(
         reference_data = self.load_reference_data(
             rule_metadata.reference_data_config, submission_info
         )
-        entities, dc_feedback_errors_uri, successful, processing_errors_uri = self.contract.apply(working_dir, entity_locations, contract_metadata)
+        entities, dc_feedback_errors_uri, successful, processing_errors_uri = self.contract.apply(
+            working_dir, entity_locations, contract_metadata
+        )
         if not successful:
             return entities, dc_feedback_errors_uri, successful, processing_errors_uri
 
@@ -172,7 +168,8 @@ def apply(
         # TODO: Handle entity manager creation errors.
         entity_manager = EntityManager(entities, reference_data)
         # TODO: Add stage success to 'apply_rules'
-        # TODO: In case of large errors in business rules, write messages to jsonl file and return uri to errors
+        # TODO: In case of large errors in business rules, write messages to jsonl file
+        # TODO: and return uri to errors
         _ = self.step_implementations.apply_rules(entity_manager, rule_metadata)
 
         for entity_name, entity in entity_manager.entities.items():
@@ -196,7 +193,9 @@ def process(
             working_dir, entity_locations, contract_metadata, rule_metadata, submission_info
         )
         if successful:
-            parquet_locations = self.write_entities_to_parquet(entities, joinuri(working_dir, "outputs"))
+            parquet_locations = self.write_entities_to_parquet(
+                entities, joinuri(working_dir, "outputs")
+            )
         else:
             parquet_locations = {}
         return parquet_locations, feedback_errors_uri, processing_errors_uri
@@ -234,6 +233,8 @@ def process_legacy(
             return entities, errors_uri  # type: ignore
 
         return (
-            self.convert_entities_to_spark(entities, joinuri(working_dir, "outputs"), _emit_deprecation_warning=False),
+            self.convert_entities_to_spark(
+                entities, joinuri(working_dir, "outputs"), _emit_deprecation_warning=False
+            ),
             errors_uri,
         )

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

@@ -9,6 +9,11 @@
 from pydantic import BaseModel
 from typing_extensions import Protocol
 
+from dve.common.error_utils import (
+    dump_processing_errors,
+    get_feedback_errors_uri,
+    get_processing_errors_uri,
+)
 from dve.core_engine.backends.base.core import get_entity_type
 from dve.core_engine.backends.base.reader import BaseFileReader
 from dve.core_engine.backends.exceptions import ReaderLacksEntityTypeSupport, render_error
@@ -27,10 +32,8 @@
     Messages,
     WrapDecorator,
 )
-from dve.parser.file_handling import get_file_suffix, get_resource_exists, get_parent
-from dve.parser.file_handling.service import joinuri
+from dve.parser.file_handling import get_file_suffix, get_resource_exists
 from dve.parser.type_hints import Extension
-from dve.common.error_utils import dump_processing_errors, get_feedback_errors_uri, get_processing_errors_uri
 
 T = TypeVar("T")
 ExtensionConfig = dict[Extension, "ReaderConfig"]
@@ -362,7 +365,12 @@ def read_raw_entities(
 
     @abstractmethod
     def apply_data_contract(
-        self, working_dir: URI, entities: Entities, entity_locations: EntityLocations, contract_metadata: DataContractMetadata, key_fields: Optional[dict[str, list[str]]] = None
+        self,
+        working_dir: URI,
+        entities: Entities,
+        entity_locations: EntityLocations,
+        contract_metadata: DataContractMetadata,
+        key_fields: Optional[dict[str, list[str]]] = None,
     ) -> tuple[Entities, URI, StageSuccessful]:
         """Apply the data contract to the raw entities, returning the validated entities
         and any messages.
@@ -373,7 +381,11 @@ def apply_data_contract(
         raise NotImplementedError()
 
     def apply(
-        self, working_dir: URI, entity_locations: EntityLocations, contract_metadata: DataContractMetadata, key_fields: Optional[dict[str, list[str]]] = None
+        self,
+        working_dir: URI,
+        entity_locations: EntityLocations,
+        contract_metadata: DataContractMetadata,
+        key_fields: Optional[dict[str, list[str]]] = None,
     ) -> tuple[Entities, URI, StageSuccessful, URI]:
         """Read the entities from the provided locations according to the data contract,
         and return the validated entities and any messages.

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

@@ -136,10 +136,11 @@ def _get_non_heterogenous_type(types: Sequence[type]) -> type:
         )
     return type_list[0]
 
+
 def check_if_parquet_file(file_location: URI) -> bool:
     """Check if a file path is valid parquet"""
     try:
         pq.ParquetFile(file_location)
         return True
     except (pyarrow.ArrowInvalid, pyarrow.ArrowIOError):
-        return False
+        return False