Modernize - phase 4 - types update

bbconf/__init__.py

@@ -3,6 +3,7 @@
 import coloredlogs
 
 from bbconf.bbagent import BedBaseAgent
+
 from .const import PKG_NAME
 
 __all__ = ["BedBaseAgent"]

bbconf/bbagent.py

@@ -2,11 +2,10 @@
 import statistics
 from functools import cached_property
 from pathlib import Path
-from typing import Dict, List, Union
 
 import numpy as np
-from sqlalchemy.orm import Session
 from sqlalchemy.engine import ScalarResult
+from sqlalchemy.orm import Session
 from sqlalchemy.sql import and_, distinct, func, or_, select
 
 from bbconf.config_parser.bedbaseconfig import BedBaseConfig
@@ -18,11 +17,11 @@
     Files,
     GeoGsmStatus,
     License,
+    ReferenceGenome,
     UsageBedMeta,
     UsageBedSetMeta,
     UsageFiles,
     UsageSearch,
-    ReferenceGenome,
 )
 from bbconf.models.base_models import (
     AllFilesInfo,
@@ -43,10 +42,10 @@
 _LOGGER = logging.getLogger(PKG_NAME)
 
 
-class BedBaseAgent(object):
+class BedBaseAgent:
     def __init__(
         self,
-        config: Union[Path, str],
+        config: Path | str,
         init_ml: bool = True,
     ):
         """
@@ -113,7 +112,6 @@ def get_detailed_stats(self, concise: bool = False) -> FileStats:
         _LOGGER.info("Getting detailed statistics for all bed files")
 
         with Session(self.config.db_engine.engine) as session:
-
             bed_compliance = {
                 f[0]: f[1]
                 for f in session.execute(
@@ -339,7 +337,7 @@ def get_detailed_usage(self) -> UsageStats:
             bed_downloads=bed_downloads,
         )
 
-    def get_list_genomes(self) -> List[str]:
+    def get_list_genomes(self) -> list[str]:
         """
         Get list of genomes from the database
 
@@ -354,7 +352,7 @@ def get_list_genomes(self) -> List[str]:
             genomes = session.execute(statement).all()
         return [result[0] for result in genomes if result[0]]
 
-    def get_list_assays(self):
+    def get_list_assays(self) -> list[str]:
         """
         Get list of genomes from the database
 
@@ -371,7 +369,7 @@ def get_list_assays(self):
         return [result[0] for result in results if result[0]]
 
     @cached_property
-    def list_of_licenses(self) -> List[str]:
+    def list_of_licenses(self) -> list[str]:
         """
         Get list of licenses from the database
 
@@ -385,7 +383,6 @@ def list_of_licenses(self) -> List[str]:
     def add_usage(self, stats: UsageModel) -> None:
 
         with Session(self.config.db_engine.engine) as session:
-
             # FILES USAGE
             reported_items_files = session.scalars(
                 select(UsageFiles).where(UsageFiles.date_to > func.now())
@@ -497,7 +494,7 @@ def add_usage(self, stats: UsageModel) -> None:
 
             session.commit()
 
-    def _stats_comments(self, sa_session: Session) -> Dict[str, int]:
+    def _stats_comments(self, sa_session: Session) -> dict[str, int]:
         """
         Get statistics about comments that are present in bed files.
 
@@ -551,7 +548,7 @@ def _stats_comments(self, sa_session: Session) -> Dict[str, int]:
             "header_comments": header_comments,
         }
 
-    def _stats_geo_status(self, sa_session: Session) -> Dict[str, int]:
+    def _stats_geo_status(self, sa_session: Session) -> dict[str, int]:
         """
         Get statistics about status of GEO bed file processing.
 
@@ -801,7 +798,7 @@ def _get_geo_stats(self, sa_session: Session) -> GEOStatistics:
             ),
         )
 
-    def get_reference_genomes(self) -> Dict[str, str]:
+    def get_reference_genomes(self) -> dict[str, str]:
         """
         Get mapping of genome aliases to reference genome names.
 

bbconf/config_parser/bedbaseconfig.py

@@ -3,7 +3,7 @@
 import os
 import warnings
 from pathlib import Path
-from typing import List, Literal, Union
+from typing import Literal
 
 import boto3
 import joblib
@@ -12,6 +12,7 @@
 import s3fs
 import yacman
 import zarr
+from botocore.client import BaseClient
 from botocore.exceptions import BotoCoreError, EndpointConnectionError
 from fastembed import TextEmbedding
 from geniml.region2vec.main import Region2VecExModel
@@ -47,12 +48,12 @@
 _LOGGER = logging.getLogger(PKG_NAME)
 
 
-class BedBaseConfig(object):
+class BedBaseConfig:
     """
     Class to handle BEDbase configuration file and create objects for different modules.
     """
 
-    def __init__(self, config: Union[Path, str], init_ml: bool = True):
+    def __init__(self, config: Path | str, init_ml: bool = True):
         """
         Initialize BedBaseConfig object
 
@@ -73,21 +74,20 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True):
             init_ml = False
 
         if init_ml:
-
             self.dense_encoder: TextEmbedding = self._init_dense_encoder()
-            self.sparse_encoder: Union[SparseEncoder, None] = self._init_sparce_model()
-            self.umap_encoder: Union[UMAP, None] = self._init_umap_model()
-            self.r2v_encoder: Union[Region2VecExModel, None] = self._init_r2v_encoder()
+            self.sparse_encoder: SparseEncoder | None = self._init_sparce_model()
+            self.umap_encoder: UMAP | None = self._init_umap_model()
+            self.r2v_encoder: Region2VecExModel | None = self._init_r2v_encoder()
 
             self._init_qdrant_hybrid(
                 qdrant_cl=self.qdrant_client,
                 dense_encoder=self.dense_encoder,
             )
 
-            self.qdrant_file_backend: Union[QdrantBackend, None] = (
+            self.qdrant_file_backend: QdrantBackend | None = (
                 self._init_qdrant_file_backend(qdrant_cl=self.qdrant_client)
             )  # used for bivec search
-            self._qdrant_text_backend: Union[QdrantBackend, None] = (
+            self._qdrant_text_backend: QdrantBackend | None = (
                 self._init_qdrant_text_backend(
                     qdrant_cl=self.qdrant_client,
                     dense_encoder=self.dense_encoder,
@@ -111,7 +111,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True):
             self.r2v_encoder = None
             self.b2b_search_interface = None
             self.bivec_search_interface = None
-            self.umap_encoder: Union[UMAP, None] = None
+            self.umap_encoder: UMAP | None = None
             self.sparse_encoder = None
 
         self._phc = self._init_pephubclient()
@@ -178,7 +178,7 @@ def boto3_client(self) -> boto3.client:
         return self._boto3_client
 
     @property
-    def zarr_root(self) -> Union[Z_GROUP, None]:
+    def zarr_root(self) -> Z_GROUP | None:
         """
         Get zarr root object (Group)
 
@@ -243,7 +243,7 @@ def _init_qdrant_client(self) -> QdrantClient:
 
     def _init_qdrant_file_backend(
         self, qdrant_cl: QdrantClient
-    ) -> Union[QdrantBackend, None]:
+    ) -> QdrantBackend | None:
         """
         Create qdrant client object using credentials provided in config file
 
@@ -255,7 +255,7 @@ def _init_qdrant_file_backend(
 
         if not isinstance(qdrant_cl, QdrantClient):
             _LOGGER.error(
-                f"Unable to create Qdrant bivec file collection, qdrant client is None."
+                "Unable to create Qdrant bivec file collection, qdrant client is None."
             )
             return None
 
@@ -270,7 +270,7 @@ def _init_qdrant_file_backend(
 
     def _init_qdrant_text_backend(
         self, qdrant_cl: QdrantClient, dense_encoder: TextEmbedding
-    ) -> Union[QdrantBackend, None]:
+    ) -> QdrantBackend | None:
         """
         Create qdrant client text embedding object using credentials provided in config file
 
@@ -283,12 +283,12 @@ def _init_qdrant_text_backend(
 
         if not isinstance(qdrant_cl, QdrantClient):
             _LOGGER.error(
-                f"Unable to create Qdrant bivec text collection, qdrant client is None."
+                "Unable to create Qdrant bivec text collection, qdrant client is None."
             )
             return None
         if not isinstance(dense_encoder, TextEmbedding):
             _LOGGER.error(
-                f"Unable to create Qdrant bivec text collection, dense encoder is None."
+                "Unable to create Qdrant bivec text collection, dense encoder is None."
             )
             return None
 
@@ -318,12 +318,12 @@ def _init_qdrant_hybrid(
 
         if not isinstance(qdrant_cl, QdrantClient):
             _LOGGER.error(
-                f"Unable to create Qdrant hybrid collection, qdrant client is None."
+                "Unable to create Qdrant hybrid collection, qdrant client is None."
             )
             return None
         if not isinstance(dense_encoder, TextEmbedding):
             _LOGGER.error(
-                f"Unable to create Qdrant hybrid collection, dense encoder is None."
+                "Unable to create Qdrant hybrid collection, dense encoder is None."
             )
             return None
 
@@ -385,7 +385,7 @@ def _init_bivec_interface(
         qdrant_file_backend: QdrantBackend,
         qdrant_text_backend: QdrantBackend,
         text_encoder: TextEmbedding,
-    ) -> Union[BiVectorSearchInterface, None]:
+    ) -> BiVectorSearchInterface | None:
         """
         Create BiVectorSearchInterface object using credentials provided in config file
 
@@ -409,8 +409,8 @@ def _init_bivec_interface(
     def _init_b2b_search_interface(
         self,
         qdrant_file_backend: QdrantBackend,
-        region_encoder: Union[Region2VecExModel, str],
-    ) -> Union[BED2BEDSearchInterface, None]:
+        region_encoder: Region2VecExModel | str,
+    ) -> BED2BEDSearchInterface | None:
         """
         Create Bed 2 BED search interface and return this object
 
@@ -430,7 +430,7 @@ def _init_b2b_search_interface(
             )
             return None
 
-    def _init_r2v_encoder(self) -> Union[Region2VecExModel, None]:
+    def _init_r2v_encoder(self) -> Region2VecExModel | None:
         """
         Create Region2VecExModel object using credentials provided in config file
         """
@@ -446,7 +446,7 @@ def _init_r2v_encoder(self) -> Union[Region2VecExModel, None]:
             )
             return None
 
-    def _init_dense_encoder(self) -> Union[None, TextEmbedding]:
+    def _init_dense_encoder(self) -> TextEmbedding | None:
         """
         Initialize dense model from the specified path or huggingface model hub
         """
@@ -457,7 +457,7 @@ def _init_dense_encoder(self) -> Union[None, TextEmbedding]:
         dense_encoder = TextEmbedding(self.config.path.text2vec)
         return dense_encoder
 
-    def _init_sparce_model(self) -> Union[None, SparseEncoder]:
+    def _init_sparce_model(self) -> SparseEncoder | None:
         """
         Initialize SparseEncoder model from the specified path or huggingface model hub
         """
@@ -472,7 +472,7 @@ def _init_sparce_model(self) -> Union[None, SparseEncoder]:
             return None
         return sparse_encoder
 
-    def _init_umap_model(self) -> Union[UMAP, None]:
+    def _init_umap_model(self) -> UMAP | None:
         """
         Load UMAP model from the specified path, or url
         """
@@ -486,21 +486,20 @@ def _init_umap_model(self) -> Union[UMAP, None]:
         model_path = self.config.path.umap_model
         umap_model = None
         if model_path.startswith(("http://", "https://")):
-
             try:
                 response = requests.get(model_path)
                 response.raise_for_status()
                 buffer = io.BytesIO(response.content)
                 umap_model = joblib.load(buffer)
-                print(f"UMAP model loaded from URL: {model_path}")
+                _LOGGER.info(f"UMAP model loaded from URL: {model_path}")
             except requests.RequestException as e:
                 _LOGGER.error(f"Error downloading UMAP model from URL: {e}")
                 return None
         else:
             try:
                 with open(model_path, "rb") as file:
                     umap_model = joblib.load(file)
-                print(f"UMAP model loaded from local path: {model_path}")
+                _LOGGER.info(f"UMAP model loaded from local path: {model_path}")
             except FileNotFoundError as e:
                 _LOGGER.error(f"Error loading UMAP model from local path: {e}")
                 return None
@@ -514,7 +513,7 @@ def _init_umap_model(self) -> Union[UMAP, None]:
 
     def _init_boto3_client(
         self,
-    ) -> Union[boto3.client, None]:
+    ) -> BaseClient | None:
         """
         Create Pephub client object using credentials provided in config file
 
@@ -532,7 +531,7 @@ def _init_boto3_client(
             warnings.warn(f"Error in creating boto3 client object: {e}", UserWarning)
             return None
 
-    def upload_s3(self, file_path: str, s3_path: Union[Path, str]) -> None:
+    def upload_s3(self, file_path: str, s3_path: Path | str) -> None:
         """
         Upload file to s3.
 
@@ -555,10 +554,10 @@ def upload_s3(self, file_path: str, s3_path: Union[Path, str]) -> None:
     def upload_files_s3(
         self,
         identifier: str,
-        files: Union[BedFiles, BedPlots, BedSetPlots],
+        files: BedFiles | BedPlots | BedSetPlots,
         base_path: str,
         type: Literal["files", "plots", "bedsets"] = "files",
-    ) -> Union[BedFiles, BedPlots, BedSetPlots]:
+    ) -> BedFiles | BedPlots | BedSetPlots:
         """
         Upload files to s3.
 
@@ -635,7 +634,7 @@ def delete_s3(self, s3_path: str) -> None:
                 "Could not delete file from s3. Connection error."
             )
 
-    def delete_files_s3(self, files: List[FileModel]) -> None:
+    def delete_files_s3(self, files: list[FileModel]) -> None:
         """
         Delete files from s3.
 
@@ -649,7 +648,7 @@ def delete_files_s3(self, files: List[FileModel]) -> None:
         return None
 
     @staticmethod
-    def _init_pephubclient() -> Union[PEPHubClient, None]:
+    def _init_pephubclient() -> PEPHubClient | None:
         """
         Create Pephub client object using credentials provided in config file
 
@@ -681,7 +680,7 @@ def get_prefixed_uri(self, postfix: str, access_id: str) -> str:
             _LOGGER.error(f"Access method {access_id} is not defined.")
             raise BadAccessMethodError(f"Access method {access_id} is not defined.")
 
-    def construct_access_method_list(self, rel_path: str) -> List[AccessMethod]:
+    def construct_access_method_list(self, rel_path: str) -> list[AccessMethod]:
         """
         Construct access method list for a given record