Type annotation improvements

lib/galaxy/app.py

@@ -608,7 +608,8 @@ def __init__(self, configure_logging=True, use_converters=True, use_display_appl
         self._register_singleton(ShortTermStorageMonitor, short_term_storage_manager)  # type: ignore[type-abstract]
 
         # Tag handler
-        self.tag_handler = self._register_singleton(GalaxyTagHandler)
+        tag_handler = GalaxyTagHandler(self.model.context)
+        self.tag_handler = self._register_singleton(GalaxyTagHandler, tag_handler)
         self.user_manager = self._register_singleton(UserManager)
         self._register_singleton(GalaxySessionManager)
         self.hda_manager = self._register_singleton(HDAManager)

lib/galaxy/app_unittest_utils/tools_support.py

@@ -130,7 +130,6 @@ def __init__(self, model_objects=None):
         self.flushed = False
         self.model_objects = model_objects or defaultdict(dict)
         self.created_objects = []
-        self.current = self
 
     def expunge_all(self):
         self.expunged_all = True

lib/galaxy/celery/base_task.py

@@ -11,9 +11,9 @@
 )
 from sqlalchemy.dialects.postgresql import insert as ps_insert
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import scoped_session
 
 from galaxy.model import CeleryUserRateLimit
-from galaxy.model.scoped_session import galaxy_scoped_session
 
 
 class GalaxyTaskBeforeStart:
@@ -45,7 +45,7 @@ class GalaxyTaskBeforeStartUserRateLimit(GalaxyTaskBeforeStart):
     def __init__(
         self,
         tasks_per_user_per_sec: float,
-        ga_scoped_session: galaxy_scoped_session,
+        ga_scoped_session: scoped_session,
     ):
         try:
             self.task_exec_countdown_secs = 1 / tasks_per_user_per_sec
@@ -68,7 +68,7 @@ def __call__(self, task: Task, task_id, args, kwargs):
 
     @abstractmethod
     def calculate_task_start_time(
-        self, user_id: int, sa_session: galaxy_scoped_session, task_interval_secs: float, now: datetime.datetime
+        self, user_id: int, sa_session: scoped_session, task_interval_secs: float, now: datetime.datetime
     ) -> datetime.datetime:
         return now
 
@@ -80,7 +80,7 @@ class GalaxyTaskBeforeStartUserRateLimitPostgres(GalaxyTaskBeforeStartUserRateLi
     """
 
     def calculate_task_start_time(
-        self, user_id: int, sa_session: galaxy_scoped_session, task_interval_secs: float, now: datetime.datetime
+        self, user_id: int, sa_session: scoped_session, task_interval_secs: float, now: datetime.datetime
     ) -> datetime.datetime:
         update_stmt = (
             update(CeleryUserRateLimit)
@@ -125,7 +125,7 @@ class GalaxyTaskBeforeStartUserRateLimitStandard(GalaxyTaskBeforeStartUserRateLi
     )
 
     def calculate_task_start_time(
-        self, user_id: int, sa_session: galaxy_scoped_session, task_interval_secs: float, now: datetime.datetime
+        self, user_id: int, sa_session: scoped_session, task_interval_secs: float, now: datetime.datetime
     ) -> datetime.datetime:
         last_scheduled_time = None
         last_scheduled_time = sa_session.scalars(self._select_stmt, {"userid": user_id}).first()

lib/galaxy/managers/context.py

@@ -63,7 +63,6 @@
     User,
 )
 from galaxy.model.base import ModelMapping
-from galaxy.model.scoped_session import galaxy_scoped_session
 from galaxy.model.tags import GalaxyTagHandlerSession
 from galaxy.schema.tasks import RequestUser
 from galaxy.security.idencoding import IdEncodingHelper
@@ -155,10 +154,10 @@ def log_event(self, message, tool_id=None, **kwargs):
             self.sa_session.commit()
 
     @property
-    def sa_session(self) -> galaxy_scoped_session:
+    def sa_session(self):
         """Provide access to Galaxy's SQLAlchemy session.
 
-        :rtype: galaxy.model.scoped_session.galaxy_scoped_session
+        :rtype: sqlalchemy.orm.scoped_session
         """
         return self.app.model.session
 

lib/galaxy/managers/histories.py

@@ -14,6 +14,7 @@
     Optional,
     Set,
     Tuple,
+    TYPE_CHECKING,
     Union,
 )
 
@@ -85,6 +86,9 @@
     RawTextTerm,
 )
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import ScalarResult
+
 log = logging.getLogger(__name__)
 
 INDEX_SEARCH_FILTERS = {
@@ -95,7 +99,7 @@
 }
 
 
-class HistoryManager(sharable.SharableModelManager, deletable.PurgableManagerMixin, SortableManager):
+class HistoryManager(sharable.SharableModelManager[model.History], deletable.PurgableManagerMixin, SortableManager):
     model_class = model.History
     foreign_key_name = "history"
     user_share_model = model.HistoryUserShareAssociation
@@ -120,7 +124,7 @@ def __init__(
 
     def index_query(
         self, trans: ProvidesUserContext, payload: HistoryIndexQueryPayload, include_total_count: bool = False
-    ) -> Tuple[List[model.History], int]:
+    ) -> Tuple["ScalarResult[model.History]", Union[int, None]]:
         show_deleted = False
         show_own = payload.show_own
         show_published = payload.show_published
@@ -234,7 +238,7 @@ def p_tag_filter(term_text: str, quoted: bool):
             stmt = stmt.limit(payload.limit)
         if payload.offset is not None:
             stmt = stmt.offset(payload.offset)
-        return trans.sa_session.scalars(stmt), total_matches  # type:ignore[return-value]
+        return trans.sa_session.scalars(stmt), total_matches
 
     # .... sharable
     # overriding to handle anonymous users' current histories in both cases

lib/galaxy/managers/history_contents.py

@@ -168,8 +168,8 @@ def state_counts(self, history):
         statement: Select = (
             select(sql.column("state"), func.count()).select_from(contents_subquery).group_by(sql.column("state"))
         )
-        counts = self.app.model.session().execute(statement).fetchall()
-        return dict(counts)
+        counts = self.app.model.session.execute(statement).fetchall()
+        return dict(counts)  # type:ignore[arg-type]
 
     def active_counts(self, history):
         """

lib/galaxy/managers/job_connections.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 from sqlalchemy import (
     literal,
     union,
@@ -9,13 +11,15 @@
 
 from galaxy import model
 from galaxy.managers.base import get_class
-from galaxy.model.scoped_session import galaxy_scoped_session
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import scoped_session
 
 
 class JobConnectionsManager:
     """Get connections graph of inputs and outputs for given item"""
 
-    def __init__(self, sa_session: galaxy_scoped_session):
+    def __init__(self, sa_session: "scoped_session"):
         self.sa_session = sa_session
 
     def get_connections_graph(self, id: int, src: str):

lib/galaxy/managers/model_stores.py

@@ -317,6 +317,7 @@ def import_model_store(self, request: ImportModelStoreTaskRequest):
 
     def _build_user_context(self, user_id: int):
         user = self._user_manager.by_id(user_id)
+        assert user is not None
         user_context = ModelStoreUserContext(self._app, user)
         return user_context
 

lib/galaxy/managers/pages.py

@@ -14,9 +14,10 @@
     Callable,
     Optional,
     Tuple,
+    TYPE_CHECKING,
+    Union,
 )
 
-import sqlalchemy
 from sqlalchemy import (
     desc,
     false,
@@ -76,6 +77,9 @@
     RawTextTerm,
 )
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import ScalarResult
+
 log = logging.getLogger(__name__)
 
 # Copied from https://github.com/kurtmckee/feedparser
@@ -121,7 +125,7 @@
 }
 
 
-class PageManager(sharable.SharableModelManager, UsesAnnotations):
+class PageManager(sharable.SharableModelManager[model.Page], UsesAnnotations):
     """Provides operations for managing a Page."""
 
     model_class = model.Page
@@ -139,7 +143,7 @@ def __init__(self, app: MinimalManagerApp):
 
     def index_query(
         self, trans: ProvidesUserContext, payload: PageIndexQueryPayload, include_total_count: bool = False
-    ) -> Tuple[sqlalchemy.engine.Result, int]:
+    ) -> Tuple["ScalarResult[model.Page]", Union[int, None]]:
         show_deleted = payload.deleted
         show_own = payload.show_own
         show_published = payload.show_published
@@ -242,7 +246,7 @@ def p_tag_filter(term_text: str, quoted: bool):
             stmt = stmt.limit(payload.limit)
         if payload.offset is not None:
             stmt = stmt.offset(payload.offset)
-        return trans.sa_session.scalars(stmt), total_matches  # type:ignore[return-value]
+        return trans.sa_session.scalars(stmt), total_matches
 
     def create_page(self, trans, payload: CreatePagePayload):
         user = trans.get_user()
@@ -269,7 +273,7 @@ def create_page(self, trans, payload: CreatePagePayload):
         content = self.rewrite_content_for_import(trans, content, content_format)
 
         # Create the new stored page
-        page = trans.app.model.Page()
+        page = model.Page()
         page.title = payload.title
         page.slug = payload.slug
         if (page_annotation := payload.annotation) is not None:
@@ -278,7 +282,7 @@ def create_page(self, trans, payload: CreatePagePayload):
 
         page.user = user
         # And the first (empty) page revision
-        page_revision = trans.app.model.PageRevision()
+        page_revision = model.PageRevision()
         page_revision.title = payload.title
         page_revision.page = page
         page.latest_revision = page_revision
@@ -310,7 +314,7 @@ def save_new_revision(self, trans, page, payload):
             content_format = page.latest_revision.content_format
         content = self.rewrite_content_for_import(trans, content, content_format=content_format)
 
-        page_revision = trans.app.model.PageRevision()
+        page_revision = model.PageRevision()
         page_revision.title = title
         page_revision.page = page
         page.latest_revision = page_revision

lib/galaxy/managers/roles.py

@@ -128,7 +128,7 @@ def purge(self, trans: ProvidesUserContext, role: model.Role) -> model.Role:
             raise RequestParameterInvalidException(f"Role '{role.name}' has not been deleted, so it cannot be purged.")
         # Delete UserRoleAssociations
         for ura in role.users:
-            user = sa_session.get(trans.app.model.User, ura.user_id)
+            user = sa_session.get(model.User, ura.user_id)
             assert user
             # Delete DefaultUserPermissions for associated users
             for dup in user.default_permissions:

lib/galaxy/managers/secured.py

@@ -33,14 +33,14 @@ class AccessibleManagerMixin:
     def by_id(self, id: int): ...
 
     # don't want to override by_id since consumers will also want to fetch w/o any security checks
-    def is_accessible(self, item, user: model.User, **kwargs: Any) -> bool:
+    def is_accessible(self, item, user: Optional[model.User], **kwargs: Any) -> bool:
         """
         Return True if the item accessible to user.
         """
         # override in subclasses
         raise exceptions.NotImplemented("Abstract interface Method")
 
-    def get_accessible(self, id: int, user: model.User, **kwargs: Any):
+    def get_accessible(self, id: int, user: Optional[model.User], **kwargs: Any):
         """
         Return the item with the given id if it's accessible to user,
         otherwise raise an error.
@@ -50,7 +50,7 @@ def get_accessible(self, id: int, user: model.User, **kwargs: Any):
         item = self.by_id(id)
         return self.error_unless_accessible(item, user, **kwargs)
 
-    def error_unless_accessible(self, item: "Query", user, **kwargs):
+    def error_unless_accessible(self, item: "Query", user: Optional[model.User], **kwargs):
         """
         Raise an error if the item is NOT accessible to user, otherwise return the item.
 

lib/galaxy/managers/sharable.py

@@ -17,6 +17,7 @@
     Optional,
     Set,
     Type,
+    TypeVar,
 )
 
 from slugify import slugify
@@ -54,10 +55,12 @@
 from galaxy.util.hash_util import md5_hash_str
 
 log = logging.getLogger(__name__)
+# Only model classes that have `users_shared_with` field
+U = TypeVar("U", model.History, model.Page, model.StoredWorkflow, model.Visualization)
 
 
 class SharableModelManager(
-    base.ModelManager,
+    base.ModelManager[U],
     secured.OwnableManagerMixin,
     secured.AccessibleManagerMixin,
     annotatable.AnnotatableManagerMixin,

lib/galaxy/managers/users.py

@@ -278,7 +278,7 @@ def _error_on_duplicate_email(self, email: str) -> None:
         if self.by_email(email) is not None:
             raise exceptions.Conflict("Email must be unique", email=email)
 
-    def by_id(self, user_id: int) -> model.User:
+    def by_id(self, user_id: int) -> Optional[model.User]:
         return self.app.model.session.get(self.model_class, user_id)
 
     # ---- filters

lib/galaxy/managers/visualizations.py

@@ -8,8 +8,9 @@
 import logging
 from typing import (
     Dict,
-    List,
     Tuple,
+    TYPE_CHECKING,
+    Union,
 )
 
 from sqlalchemy import (
@@ -44,6 +45,9 @@
     RawTextTerm,
 )
 
+if TYPE_CHECKING:
+    from sqlalchemy.engine import ScalarResult
+
 log = logging.getLogger(__name__)
 
 
@@ -59,7 +63,7 @@
 }
 
 
-class VisualizationManager(sharable.SharableModelManager):
+class VisualizationManager(sharable.SharableModelManager[model.Visualization]):
     """
     Handle operations outside and between visualizations and other models.
     """
@@ -76,7 +80,7 @@ class VisualizationManager(sharable.SharableModelManager):
 
     def index_query(
         self, trans: ProvidesUserContext, payload: VisualizationIndexQueryPayload, include_total_count: bool = False
-    ) -> Tuple[List[model.Visualization], int]:
+    ) -> Tuple["ScalarResult[model.Visualization]", Union[int, None]]:
         show_deleted = payload.deleted
         show_own = payload.show_own
         show_published = payload.show_published
@@ -171,7 +175,7 @@ def p_tag_filter(term_text: str, quoted: bool):
             stmt = stmt.limit(payload.limit)
         if payload.offset is not None:
             stmt = stmt.offset(payload.offset)
-        return trans.sa_session.scalars(stmt), total_matches  # type:ignore[return-value]
+        return trans.sa_session.scalars(stmt), total_matches
 
 
 class VisualizationSerializer(sharable.SharableModelSerializer):