Improve type annotations

Author: nsoranzo

lib/galaxy/managers/jobs.py

@@ -12,6 +12,7 @@
     cast,
     Optional,
     TYPE_CHECKING,
+    TypeVar,
     Union,
 )
 
@@ -117,7 +118,11 @@
 from galaxy.work.context import WorkRequestContext
 
 if TYPE_CHECKING:
-    from sqlalchemy.sql.expression import Select
+    from sqlalchemy.sql.expression import (
+        ColumnElement,
+        Label,
+        Select,
+    )
 
 log = logging.getLogger(__name__)
 
@@ -133,7 +138,7 @@ class JobLock(BaseModel):
     active: bool = Field(title="Job lock status", description="If active, jobs will not dispatch")
 
 
-def get_path_key(path_tuple):
+def get_path_key(path_tuple: tuple):
     path_key = ""
     tuple_elements = len(path_tuple)
     for i, p in enumerate(path_tuple):
@@ -153,12 +158,15 @@ def get_path_key(path_tuple):
 
 
 def safe_label_or_none(label: str) -> Optional[str]:
-    if label and len(label) > 63:
+    if len(label) > 63:
         return None
     return label
 
 
-def safe_aliased(model_class, name=None):
+T = TypeVar("T")
+
+
+def safe_aliased(model_class: type[T], name: str) -> type[T]:
     """Create an aliased model class with a unique name."""
     return aliased(model_class, name=safe_label_or_none(name))
 
@@ -476,11 +484,11 @@ def by_tool_input(
         job_state: Optional[JobStatesT] = (Job.states.OK,),
         history_id: Union[int, None] = None,
         require_name_match: bool = True,
-    ):
+    ) -> Union[Job, None]:
         """Search for jobs producing same results using the 'inputs' part of a tool POST."""
-        input_data = defaultdict(list)
+        input_data: dict[Any, list[dict[str, Any]]] = defaultdict(list)
 
-        def populate_input_data_input_id(path, key, value):
+        def populate_input_data_input_id(path: tuple, key, value) -> tuple[Any, Any]:
             """Traverses expanded incoming using remap and collects input_ids and input_data."""
             if key == "id":
                 path_key = get_path_key(path[:-2])
@@ -528,13 +536,13 @@ def __search(
         tool_id: str,
         tool_version: Optional[str],
         user: model.User,
-        input_data,
+        input_data: dict[Any, list[dict[str, Any]]],
         job_state: Optional[JobStatesT],
         param_dump: ToolStateDumpedToJsonInternalT,
         wildcard_param_dump=None,
         history_id: Union[int, None] = None,
         require_name_match: bool = True,
-    ):
+    ) -> Union[Job, None]:
         search_timer = ExecutionTimer()
 
         def replace_dataset_ids(path, key, value):
@@ -554,15 +562,15 @@ def replace_dataset_ids(path, key, value):
 
         stmt = select(model.Job.id.label("job_id"))
 
-        data_conditions: list = []
+        data_conditions: list[ColumnElement[bool]] = []
 
         # We now build the stmt filters that relate to the input datasets
         # that this job uses. We keep track of the requested dataset id in `requested_ids`,
         # the type (hda, hdca or lda) in `data_types`
         # and the ids that have been used in the job that has already been run in `used_ids`.
         requested_ids = []
         data_types = []
-        used_ids: list = []
+        used_ids: list[Label[int]] = []
         for k, input_list in input_data.items():
             # k will be matched against the JobParameter.name column. This can be prefixed depending on whether
             # the input is in a repeat, or not (section and conditional)
@@ -751,7 +759,7 @@ def _filter_jobs(
 
         return stmt
 
-    def _exclude_jobs_with_deleted_outputs(self, stmt):
+    def _exclude_jobs_with_deleted_outputs(self, stmt: "Select[tuple[int]]") -> "Select":
         subquery_alias = stmt.subquery("filtered_jobs_subquery")
         outer_select_columns = [subquery_alias.c[col.name] for col in stmt.selected_columns]
         outer_stmt = select(*outer_select_columns).select_from(subquery_alias)
@@ -796,14 +804,14 @@ def _exclude_jobs_with_deleted_outputs(self, stmt):
     def _build_stmt_for_hda(
         self,
         stmt: "Select[tuple[int]]",
-        data_conditions: list,
-        used_ids: list,
+        data_conditions: list["ColumnElement[bool]"],
+        used_ids: list["Label[int]"],
         k,
         v,
         identifier,
         value_index: int,
         require_name_match: bool = True,
-    ):
+    ) -> "Select[tuple[int]]":
         a = aliased(model.JobToInputDatasetAssociation)
         b = aliased(model.HistoryDatasetAssociation)
         c = aliased(model.HistoryDatasetAssociation)
@@ -859,7 +867,15 @@ def _build_stmt_for_hda(
         )
         return stmt
 
-    def _build_stmt_for_ldda(self, stmt, data_conditions, used_ids, k, v, value_index):
+    def _build_stmt_for_ldda(
+        self,
+        stmt: "Select[tuple[int]]",
+        data_conditions: list["ColumnElement[bool]"],
+        used_ids: list["Label[int]"],
+        k,
+        v,
+        value_index: int,
+    ) -> "Select[tuple[int]]":
         a = aliased(model.JobToInputLibraryDatasetAssociation)
         label = safe_label_or_none(f"{k}_{value_index}")
         labeled_col = a.ldda_id.label(label)
@@ -876,8 +892,15 @@ def agg_expression(self, column):
             return func.array_agg(column, order_by=column)
 
     def _build_stmt_for_hdca(
-        self, stmt, data_conditions, used_ids, k, v, user_id, value_index, require_name_match=True
-    ):
+        self,
+        stmt: "Select[tuple[int]]",
+        data_conditions: list["ColumnElement[bool]"],
+        used_ids: list["Label[int]"],
+        k,
+        v,
+        user_id: int,
+        value_index: int,
+    ) -> "Select[tuple[int]]":
         # Strategy for efficiently finding equivalent HDCAs:
         # 1. Determine the structural depth of the target HDCA by its collection_type.
         # 2. For the target HDCA (identified by 'v'):
@@ -1112,7 +1135,16 @@ def _build_stmt_for_hdca(
         data_conditions.append(a.name == k)
         return stmt
 
-    def _build_stmt_for_dce(self, stmt, data_conditions, used_ids, k, v, user_id, value_index):
+    def _build_stmt_for_dce(
+        self,
+        stmt: "Select[tuple[int]]",
+        data_conditions: list["ColumnElement[bool]"],
+        used_ids: list["Label[int]"],
+        k,
+        v,
+        user_id: int,
+        value_index: int,
+    ) -> "Select[tuple[int]]":
         dce_root_target = self.sa_session.get_one(model.DatasetCollectionElement, v)
 
         # Determine if the target DCE points to an HDA or a child collection

lib/galaxy_test/api/test_tool_execute.py

@@ -755,7 +755,7 @@ def _run_deferred(
 
 
 @requires_tool_id("cat|cat1")
-def test_deferred_with_cached_input(required_tool: RequiredTool, target_history: TargetHistory):
+def test_deferred_with_cached_input(required_tool: RequiredTool, target_history: TargetHistory) -> None:
     # Basic deferred dataset
     _run_deferred(required_tool, target_history)
     # Should just work because input is deferred

lib/galaxy_test/base/api.py

@@ -3,6 +3,7 @@
 from typing import (
     Any,
     Optional,
+    TYPE_CHECKING,
 )
 from urllib.parse import (
     urlencode,
@@ -30,6 +31,9 @@
 )
 from .interactor import TestCaseGalaxyInteractor as BaseInteractor
 
+if TYPE_CHECKING:
+    from requests import Response
+
 CONFIG_PREFIXES = ["GALAXY_TEST_CONFIG_", "GALAXY_CONFIG_OVERRIDE_", "GALAXY_CONFIG_"]
 CELERY_BROKER = get_from_env("CELERY_BROKER", CONFIG_PREFIXES, "memory://")
 CELERY_BACKEND = get_from_env("CELERY_BACKEND", CONFIG_PREFIXES, "rpc://localhost")
@@ -198,7 +202,7 @@ def _patch(self, *args, **kwds):
     def _assert_status_code_is_ok(self, response):
         assert_status_code_is_ok(response)
 
-    def _assert_status_code_is(self, response, expected_status_code):
+    def _assert_status_code_is(self, response: "Response", expected_status_code: int) -> None:
         assert_status_code_is(response, expected_status_code)
 
     def _assert_has_keys(self, response, *keys):