refactor(package): Streamline typing imports and fix DockerMount type annotation. (#1186)

Author: haiqi96

components/clp-package-utils/clp_package_utils/scripts/archive_manager.py

@@ -2,8 +2,8 @@
 import logging
 import subprocess
 import sys
-import typing
 from pathlib import Path
+from typing import Final, List, Optional
 
 from clp_py_utils.clp_config import (
     CLP_DEFAULT_DATASET_NAME,
@@ -14,7 +14,7 @@
 from clp_package_utils.general import (
     CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH,
     CLPConfig,
-    CLPDockerMounts,
+    DockerMount,
     dump_container_config,
     generate_container_config,
     generate_container_name,
@@ -26,18 +26,18 @@
 )
 
 # Command/Argument Constants
-FIND_COMMAND: typing.Final[str] = "find"
-DEL_COMMAND: typing.Final[str] = "del"
-DEL_BY_IDS_SUBCOMMAND: typing.Final[str] = "by-ids"
-DEL_BY_FILTER_SUBCOMMAND: typing.Final[str] = "by-filter"
-BEGIN_TS_ARG: typing.Final[str] = "--begin-ts"
-END_TS_ARG: typing.Final[str] = "--end-ts"
-DRY_RUN_ARG: typing.Final[str] = "--dry-run"
+FIND_COMMAND: Final[str] = "find"
+DEL_COMMAND: Final[str] = "del"
+DEL_BY_IDS_SUBCOMMAND: Final[str] = "by-ids"
+DEL_BY_FILTER_SUBCOMMAND: Final[str] = "by-filter"
+BEGIN_TS_ARG: Final[str] = "--begin-ts"
+END_TS_ARG: Final[str] = "--end-ts"
+DRY_RUN_ARG: Final[str] = "--dry-run"
 
 logger: logging.Logger = logging.getLogger(__file__)
 
 
-def _validate_timestamps(begin_ts: int, end_ts: typing.Optional[int]) -> bool:
+def _validate_timestamps(begin_ts: int, end_ts: Optional[int]) -> bool:
     if begin_ts < 0:
         logger.error("begin-ts must be non-negative.")
         return False
@@ -50,7 +50,7 @@ def _validate_timestamps(begin_ts: int, end_ts: typing.Optional[int]) -> bool:
     return True
 
 
-def main(argv: typing.List[str]) -> int:
+def main(argv: List[str]) -> int:
     clp_home: Path = get_clp_home()
     default_config_file_path: Path = clp_home / CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH
 
@@ -159,8 +159,8 @@ def main(argv: typing.List[str]) -> int:
     else:
         logger.setLevel(logging.INFO)
 
-    begin_timestamp: typing.Optional[int]
-    end_timestamp: typing.Optional[int]
+    begin_timestamp: Optional[int] = None
+    end_timestamp: Optional[int] = None
     subcommand: str = parsed_args.subcommand
 
     # Validate and load config file
@@ -215,17 +215,17 @@ def main(argv: typing.List[str]) -> int:
         container_clp_config, clp_config, container_name
     )
 
-    necessary_mounts: typing.List[CLPDockerMounts] = [
+    necessary_mounts: List[Optional[DockerMount]] = [
         mounts.clp_home,
         mounts.logs_dir,
         mounts.archives_output_dir,
     ]
-    container_start_cmd: typing.List[str] = generate_container_start_cmd(
+    container_start_cmd: List[str] = generate_container_start_cmd(
         container_name, necessary_mounts, clp_config.execution_container
     )
 
     # fmt: off
-    archive_manager_cmd: typing.List[str] = [
+    archive_manager_cmd: List[str] = [
         "python3",
         "-m", "clp_package_utils.scripts.native.archive_manager",
         "--config", str(generated_config_path_on_container),
@@ -264,7 +264,7 @@ def main(argv: typing.List[str]) -> int:
         logger.error(f"Unsupported subcommand: `{subcommand}`.")
         return -1
 
-    cmd: typing.List[str] = container_start_cmd + archive_manager_cmd
+    cmd: List[str] = container_start_cmd + archive_manager_cmd
 
     proc = subprocess.run(cmd)
     ret_code = proc.returncode

components/clp-package-utils/clp_package_utils/scripts/native/archive_manager.py

@@ -2,10 +2,10 @@
 import logging
 import shutil
 import sys
-import typing
 from abc import ABC, abstractmethod
 from contextlib import closing
 from pathlib import Path
+from typing import Any, List, Optional
 
 from clp_py_utils.clp_config import Database
 from clp_py_utils.clp_metadata_db_utils import (
@@ -35,10 +35,10 @@
 
 
 class DeleteHandler(ABC):
-    def __init__(self, query_params: typing.List[str]):
-        self._params: typing.List[str] = query_params
+    def __init__(self, query_params: List[str]):
+        self._params: List[str] = query_params
 
-    def get_params(self) -> typing.List[str]:
+    def get_params(self) -> List[str]:
         return self._params
 
     @abstractmethod
@@ -48,7 +48,7 @@ def get_criteria(self) -> str: ...
     def get_not_found_message(self) -> str: ...
 
     @abstractmethod
-    def validate_results(self, archive_ids: typing.List[str]) -> None: ...
+    def validate_results(self, archive_ids: List[str]) -> None: ...
 
 
 class FilterDeleteHandler(DeleteHandler):
@@ -58,7 +58,7 @@ def get_criteria(self) -> str:
     def get_not_found_message(self) -> str:
         return "No archives found within the specified time range."
 
-    def validate_results(self, archive_ids: typing.List[str]) -> None:
+    def validate_results(self, archive_ids: List[str]) -> None:
         pass
 
 
@@ -70,7 +70,7 @@ def get_criteria(self) -> str:
     def get_not_found_message(self) -> str:
         return "No archives found with matching IDs."
 
-    def validate_results(self, archive_ids: typing.List[str]) -> None:
+    def validate_results(self, archive_ids: List[str]) -> None:
         not_found_ids: set[str] = set(self._params) - set(archive_ids)
         if not_found_ids:
             logger.warning(
@@ -81,7 +81,7 @@ def validate_results(self, archive_ids: typing.List[str]) -> None:
             )
 
 
-def main(argv: typing.List[str]) -> int:
+def main(argv: List[str]) -> int:
     clp_home: Path = get_clp_home()
     default_config_file_path: Path = clp_home / CLP_DEFAULT_CONFIG_FILE_RELATIVE_PATH
 
@@ -253,9 +253,9 @@ def main(argv: typing.List[str]) -> int:
 def _find_archives(
     archives_dir: Path,
     database_config: Database,
-    dataset: typing.Optional[str],
+    dataset: Optional[str],
     begin_ts: int,
-    end_ts: int = typing.Optional[int],
+    end_ts: int,
 ) -> int:
     """
     Lists all archive IDs, if begin_ts and end_ts are provided, only lists archives where
@@ -267,20 +267,20 @@ def _find_archives(
     :param end_ts:
     :return: 0 on success, 1 on failure.
     """
-    archive_ids: typing.List[str]
+    archive_ids: List[str]
     dataset_specific_message = f" of dataset `{dataset}`" if dataset is not None else ""
     logger.info(f"Starting to find archives{dataset_specific_message} from the database.")
     try:
         sql_adapter: SQL_Adapter = SQL_Adapter(database_config)
-        clp_db_connection_params: dict[str, any] = (
+        clp_db_connection_params: dict[str, Any] = (
             database_config.get_clp_connection_params_and_type(True)
         )
         table_prefix: str = clp_db_connection_params["table_prefix"]
 
         with closing(sql_adapter.create_connection(True)) as db_conn, closing(
             db_conn.cursor(dictionary=True)
         ) as db_cursor:
-            query_params: typing.List[int] = [begin_ts]
+            query_params: List[int] = [begin_ts]
             query: str = (
                 f"""
                 SELECT id FROM `{get_archives_table_name(table_prefix, dataset)}`
@@ -294,7 +294,7 @@ def _find_archives(
             db_cursor.execute(query, query_params)
             results = db_cursor.fetchall()
 
-            archive_ids: typing.List[str] = [result["id"] for result in results]
+            archive_ids: List[str] = [result["id"] for result in results]
             if 0 == len(archive_ids):
                 logger.info("No archives found within specified time range.")
                 return 0
@@ -318,7 +318,7 @@ def _find_archives(
 def _delete_archives(
     archives_dir: Path,
     database_config: Database,
-    dataset: typing.Optional[str],
+    dataset: Optional[str],
     delete_handler: DeleteHandler,
     dry_run: bool = False,
 ) -> int:
@@ -333,11 +333,11 @@ def _delete_archives(
     :return: 0 on success, -1 otherwise.
     """
 
-    archive_ids: typing.List[str]
+    archive_ids: List[str]
     dataset_specific_message = f" of dataset `{dataset}`" if dataset is not None else ""
     logger.info(f"Starting to delete archives{dataset_specific_message} from the database.")
     sql_adapter: SQL_Adapter = SQL_Adapter(database_config)
-    clp_db_connection_params: dict[str, any] = database_config.get_clp_connection_params_and_type(
+    clp_db_connection_params: dict[str, Any] = database_config.get_clp_connection_params_and_type(
         True
     )
     table_prefix = clp_db_connection_params["table_prefix"]
@@ -350,7 +350,7 @@ def _delete_archives(
                 logger.info("Running in dry-run mode.")
 
             query_criteria: str = delete_handler.get_criteria()
-            query_params: typing.List[str] = delete_handler.get_params()
+            query_params: List[str] = delete_handler.get_params()
 
             db_cursor.execute(
                 f"""
@@ -365,7 +365,7 @@ def _delete_archives(
                 logger.info(delete_handler.get_not_found_message())
                 return 0
 
-            archive_ids: typing.List[str] = [result["id"] for result in results]
+            archive_ids: List[str] = [result["id"] for result in results]
             delete_handler.validate_results(archive_ids)
 
             delete_archives_from_metadata_db(db_cursor, archive_ids, table_prefix, dataset)