sibi-dst 2025.1.8__tar.gz → 2025.1.10__tar.gz

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sibi-dst
-Version: 2025.1.8
+Version: 2025.1.10
 Summary: Data Science Toolkit
 Author: Luis Valverde
 Author-email: lvalverdeb@gmail.com

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sibi-dst"
-version = "2025.1.8"
+version = "2025.1.10"
 description = "Data Science Toolkit"
 authors = ["Luis Valverde <lvalverdeb@gmail.com>"]
 readme = "README.md"

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/df_helper/_artifact_updater_multi_wrapper.py

@@ -49,6 +49,7 @@ class ArtifactUpdaterMultiWrapperThreaded(ManagedResource):
         self.completion_times: Dict[str, float] = {}
         self.failed: List[str] = []
         self.original_classes: List[Type] = []
+        self.logger.info("ArtifactUpdaterMultiWrapperThreaded initialized")
 
     def get_artifact_classes(self, data_type: str) -> List[Type]:
         """Retrieve artifact classes by data type."""
@@ -270,6 +271,7 @@ class ArtifactUpdaterMultiWrapperAsync(ManagedResource):
         self.completion_times: Dict[str, float] = {}
         self.failed: List[str] = []
         self.original_classes: List[Type] = []
+        self.logger.info("ArtifactUpdaterMultiWrapperAsync initialized")
 
     def get_artifact_classes(self, data_type: str) -> List[Type]:
         """

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/df_helper/_df_helper.py

@@ -28,6 +28,7 @@ class BaseBackend:
         self.logger = helper.logger
         self.debug = helper.debug
         self.total_records = helper.total_records  # no records loaded yet
+        self._entered = helper._entered  # Track if the helper is used in a context manager
 
     def load(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         """Synchronous data loading method. Must be implemented by sync backends."""
@@ -67,7 +68,7 @@ class ParquetBackend(BaseBackend):
             df = self.helper.backend_parquet.load_files()
             if options and df is not None:
                 df = FilterHandler('dask', logger=self.logger, debug=False).apply_filters(df, filters=options)
-            self.total_records = len(df)
+            self.total_records = len(df) or -1  # If df is empty, set total_records to -1
             return self.total_records, df
         except Exception as e:
             self.total_records = -1  # Reset total_records on failure
@@ -105,6 +106,12 @@ class DfHelper(ManagedResource):
         'http': HttpBackend,
     }
 
+    _BACKEND_ATTR_MAP = {
+        'sqlalchemy': 'backend_db_connection',
+        'parquet': 'backend_parquet',
+        'http': 'backend_http',
+    }
+
     default_config: Dict = None
 
     def __init__(self, backend='sqlalchemy', **kwargs):
@@ -140,9 +147,15 @@ class DfHelper(ManagedResource):
         super().__exit__(exc_type, exc_value, traceback)
 
     def _cleanup(self):
-        active_config = getattr(self, f"backend_{self.backend}", None)
+        attr_name = self._BACKEND_ATTR_MAP.get(self.backend)
+        if not attr_name:
+            self.logger.warning(f"No attribute mapping found for backend '{self.backend}'. Cleanup skipped.")
+            return
+        # Get the actual config object (e.g., self.backend_db_connection)
+        active_config = getattr(self, attr_name, None)
+
         if active_config and hasattr(active_config, "close"):
-            self.logger.debug(f"Closing resources for '{self.backend}' backend.")
+            self.logger.debug(f"Closing resources for '{self.backend}' backend using attribute '{attr_name}'.")
             active_config.close()
 
     def _get_config(self, model: T, kwargs: Dict[str, Any]) -> T:
@@ -156,6 +169,10 @@ class DfHelper(ManagedResource):
         self.total_records, df = self.backend_strategy.load(**options)
         df = self._process_loaded_data(df)
         df = self._post_process_df(df)
+        if not self._entered:
+            self.logger.warning(
+                "DfHelper instance was not used in a context manager; cleanup is being called manually.")
+            self._cleanup()
         return df.compute() if as_pandas else df
 
     async def aload(self, as_pandas=False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
@@ -200,7 +217,11 @@ class DfHelper(ManagedResource):
             self.logger.warning("Cannot save to parquet; DataFrame is empty.")
             return
         fs = kwargs.pop('fs', self.fs)
-        path = kwargs.pop('parquet_storage_path', self.backend_parquet.parquet_storage_path)
+        if not fs:
+            raise ValueError("Filesystem (fs) must be provided to save to parquet.")
+        path = kwargs.pop('parquet_storage_path', None)
+        if not path:
+            raise ValueError("parquet_storage_path must be provided to save to parquet.")
         writer_config = {
             'df_result': df,
             'parquet_storage_path': path,

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/df_helper/backends/parquet/_parquet_options.py

@@ -4,8 +4,8 @@ from typing import Optional, List
 
 import dask.dataframe as dd
 import fsspec
-from pydantic import BaseModel, model_validator, DirectoryPath, FilePath, ConfigDict
-
+import pandas as pd
+from pydantic import BaseModel, model_validator, ConfigDict
 from sibi_dst.utils import FilePathGenerator
 from sibi_dst.utils import Logger
 
@@ -93,7 +93,7 @@ class ParquetConfig(BaseModel):
         self.parquet_storage_path = self.parquet_storage_path.rstrip('/')
         if not self.fs.exists(self.parquet_storage_path):
             self.fs.mkdirs(self.parquet_storage_path, exist_ok=True)
-            #raise ValueError('Parquet storage path does not exist')
+            # raise ValueError('Parquet storage path does not exist')
         self.load_parquet = False
         if self.parquet_filename is not None:
             self.parquet_full_path = self.ensure_file_extension(
@@ -184,11 +184,36 @@ class ParquetConfig(BaseModel):
         :return: A Dask DataFrame containing loaded parquet file data.
         :rtype: dask.dataframe.DataFrame
         """
-        if self.load_parquet:
-            if self.parquet_folder_list:
-                return dd.read_parquet(self.parquet_folder_list, engine="pyarrow", filesystem=self.fs)
-            else:
-                return dd.read_parquet(self.parquet_full_path, engine="pyarrow", filesystem=self.fs)
+        if not self.load_parquet:
+            self.logger.warning("Parquet loading is disabled. Returning empty DataFrame.")
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+        paths_to_load = []
+        if self.parquet_folder_list:
+            # Filter out any None values from the list
+            paths_to_load = [p for p in self.parquet_folder_list if p is not None]
+        elif self.parquet_full_path:
+            # Treat the single path as a list with one item
+            paths_to_load = [self.parquet_full_path]
+
+        if not paths_to_load:
+            self.logger.warning("No valid parquet file paths were provided. Returning empty DataFrame.")
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+        try:
+            self.logger.debug(f"Attempting to load Parquet data from: {paths_to_load}")
+            return dd.read_parquet(
+                paths_to_load,
+                engine="pyarrow",
+                filesystem=self.fs,
+                exclude=["_*", ".*"]
+            )
+        except Exception as e:
+            # This robust error handling is excellent.
+            self.logger.error(f"Parquet loading failed for paths {paths_to_load}: {e}", exc_info=True)
+            self.logger.warning("Returning empty DataFrame due to loading error.")
+            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+
 
     @staticmethod
     def ensure_file_extension(filepath: str, extension: str) -> str:

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -172,7 +172,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
                 return
 
             engine_wrapper['ref_count'] -= 1
-            self.logger.debug(f"Closing config. Ref count is now {engine_wrapper['ref_count']}.")
+            self.logger.debug(f"Closing connection within engine wrapper. Ref count is now {engine_wrapper['ref_count']}.")
 
             if engine_wrapper['ref_count'] <= 0:
                 self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}")

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/utils/data_wrapper.py

@@ -153,44 +153,44 @@ class DataWrapper(ManagedResource):
             # Create a copy to avoid mutating the shared instance dictionary
             local_load_params = self.load_params.copy()
             local_load_params.update(date_filter)
-            local_class_instance = self.dataclass(**self.class_params)
-            df = local_class_instance.load(**local_load_params)
-            load_time = time.perf_counter() - load_start
-
-            if hasattr(local_class_instance, "total_records"):
-                self.logger.debug(
-                    f"Total records loaded by {local_class_instance.__class__.__name__}: {local_class_instance.total_records}")
-                if int(local_class_instance.total_records) == 0:  # If no records were loaded but not due to an error
-                    if self.mmanifest:
-                        self.mmanifest.record(
+            with self.dataclass(**self.class_params) as local_class_instance:
+                df = local_class_instance.load(**local_load_params)
+                load_time = time.perf_counter() - load_start
+
+                if hasattr(local_class_instance, "total_records"):
+                    self.logger.debug(
+                        f"Total records loaded by {local_class_instance.__class__.__name__}: {local_class_instance.total_records}")
+                    if int(local_class_instance.total_records) == 0:  # If no records were loaded but not due to an error
+                        if self.mmanifest:
+                            self.mmanifest.record(
                             full_path=path
                         )
-                    self.logger.info(f"No data found for {full_path}. Logged to missing manifest.")
-                elif int(local_class_instance.total_records) < 0:
-                    self.logger.warning(
-                        f"Negative record count ({local_class_instance.total_records}) for {full_path}. "
-                        "This may indicate an error in the data loading process."
-                    )
-                else:
-                    save_start = time.perf_counter()
-                    parquet_params ={
-                        "df_result": df,
-                        "parquet_storage_path": path,
-                        "fs": self.fs,
-                        "logger": self.logger,
-                        "debug": self.debug,
-                    }
-                    with ParquetSaver(**parquet_params) as ps:
-                        ps.save_to_parquet(self.parquet_filename, overwrite=True)
-                    save_time = time.perf_counter() - save_start
-
-                    total_time = time.perf_counter() - overall_start
-                    self.benchmarks[date] = {
-                        "load_duration": load_time,
-                        "save_duration": save_time,
-                        "total_duration": total_time
-                    }
-                    self._log_success(date, total_time, full_path)
+                        self.logger.info(f"No data found for {full_path}. Logged to missing manifest.")
+                    elif int(local_class_instance.total_records) < 0:
+                        self.logger.warning(
+                            f"Negative record count ({local_class_instance.total_records}) for {full_path}. "
+                            "This may indicate an error in the data loading process."
+                        )
+                    else:
+                        save_start = time.perf_counter()
+                        parquet_params ={
+                            "df_result": df,
+                            "parquet_storage_path": path,
+                            "fs": self.fs,
+                            "logger": self.logger,
+                            "debug": self.debug,
+                        }
+                        with ParquetSaver(**parquet_params) as ps:
+                            ps.save_to_parquet(self.parquet_filename, overwrite=True)
+                        save_time = time.perf_counter() - save_start
+
+                        total_time = time.perf_counter() - overall_start
+                        self.benchmarks[date] = {
+                            "load_duration": load_time,
+                            "save_duration": save_time,
+                            "total_duration": total_time
+                        }
+                        self._log_success(date, total_time, full_path)
         except Exception as e:
             self._log_failure(date, e)
             raise

{sibi_dst-2025.1.8 → sibi_dst-2025.1.10}/sibi_dst/utils/update_planner.py

@@ -72,9 +72,6 @@ class UpdatePlanner(ManagedResource):
             data_path: str,
             filename: str,
             description: str = "Update Planner",
-            #fs: Optional[fsspec.AbstractFileSystem] = None,
-            #filesystem_type: str = "file",
-            #filesystem_options: Optional[Dict] = None,
             reference_date: Union[str, datetime.date] = None,
             history_days_threshold: int = DEFAULT_HISTORY_DAYS_THRESHOLD,
             max_age_minutes: int = DEFAULT_MAX_AGE_MINUTES,
@@ -199,11 +196,12 @@ class UpdatePlanner(ManagedResource):
         for _, row in self.plan.iterrows():
             table.add_row(*(str(item) for item in row))
 
-        console.print(table)
+        console = Console()
+        with console.capture() as capture:
+            console.print(table)
+        plan_string = capture.get()
 
-        plan_string = console.export_text()
-
-        self.logger.info(f"Full Update Plan:\n{plan_string}")
+        self.logger.info(f"Full Update Plan:\n{plan_string.strip()}")
 
     def get_tasks_by_priority(self) -> Iterator[Tuple[int, List[datetime.date]]]:
         """Yields batches of dates to be processed, grouped and sorted by priority."""