Improve type annotation of SharedModelMapping

lib/galaxy/model/base.py

@@ -16,6 +16,7 @@
     Dict,
     List,
     Type,
+    TYPE_CHECKING,
     Union,
 )
 
@@ -28,6 +29,20 @@
 
 from galaxy.util.bunch import Bunch
 
+if TYPE_CHECKING:
+    from galaxy.model import (
+        APIKeys as GalaxyAPIKeys,
+        GalaxySession as GalaxyGalaxySession,
+        PasswordResetToken as GalaxyPasswordResetToken,
+        User as GalaxyUser,
+    )
+    from tool_shed.webapp.model import (
+        APIKeys as ToolShedAPIKeys,
+        GalaxySession as ToolShedGalaxySession,
+        PasswordResetToken as ToolShedPasswordResetToken,
+        User as ToolShedUser,
+    )
+
 log = logging.getLogger(__name__)
 
 # Create a ContextVar with mutable state, this allows sync tasks in the context
@@ -128,10 +143,10 @@ class SharedModelMapping(ModelMapping):
     a way to do app.model.<CLASS> for common code shared by the tool shed and Galaxy.
     """
 
-    User: Type
-    GalaxySession: Type
-    APIKeys: Type
-    PasswordResetToken: Type
+    User: Union[Type["GalaxyUser"], Type["ToolShedUser"]]
+    GalaxySession: Union[Type["GalaxyGalaxySession"], Type["ToolShedGalaxySession"]]
+    APIKeys: Union[Type["GalaxyAPIKeys"], Type["ToolShedAPIKeys"]]
+    PasswordResetToken: Union[Type["GalaxyPasswordResetToken"], Type["ToolShedPasswordResetToken"]]
 
 
 def versioned_objects(iter):

lib/galaxy/model/mapping.py

@@ -2,6 +2,7 @@
 from threading import local
 from typing import (
     Optional,
+    Type,
     TYPE_CHECKING,
 )
 
@@ -17,6 +18,7 @@
 from galaxy.model.triggers.update_audit_table import install as install_timestamp_triggers
 
 if TYPE_CHECKING:
+    from galaxy.model import User as GalaxyUser
     from galaxy.objectstore import BaseObjectStore
 
 log = logging.getLogger(__name__)
@@ -25,6 +27,7 @@
 
 
 class GalaxyModelMapping(SharedModelMapping):
+    User: Type["GalaxyUser"]
     security_agent: GalaxyRBACAgent
     thread_local_log: Optional[local]
 

lib/tool_shed/webapp/model/mapping.py

@@ -8,6 +8,8 @@
     Any,
     Dict,
     Optional,
+    Type,
+    TYPE_CHECKING,
 )
 
 import tool_shed.webapp.model
@@ -17,12 +19,16 @@
 from tool_shed.webapp.model import mapper_registry
 from tool_shed.webapp.security import CommunityRBACAgent
 
+if TYPE_CHECKING:
+    from tool_shed.webapp.model import User as ToolShedUser
+
 log = logging.getLogger(__name__)
 
 metadata = mapper_registry.metadata
 
 
 class ToolShedModelMapping(SharedModelMapping):
+    User: Type["ToolShedUser"]
     security_agent: CommunityRBACAgent
     shed_counter: shed_statistics.ShedCounter
     create_tables: bool