pep-484: session, instancestate, etc

Also adds some fixes to annotation-based mapping
that have come up, as well as starts to add more
pep-484 test cases

Change-Id: Ia722bbbc7967a11b23b66c8084eb61df9d233fee

1 files changed, 31 insertions, 10 deletions

diff --git a/lib/sqlalchemy/orm/unitofwork.py b/lib/sqlalchemy/orm/unitofwork.py
index da098e8c5..9ff284e73 100644
--- a/lib/sqlalchemy/orm/unitofwork.py
+++ b/lib/sqlalchemy/orm/unitofwork.py
@@ -15,6 +15,12 @@ organizes them in order of dependency, and executes.
 
 from __future__ import annotations
 
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Set
+from typing import TYPE_CHECKING
+
 from . import attributes
 from . import exc as orm_exc
 from . import util as orm_util
@@ -23,6 +29,15 @@ from .. import util
 from ..util import topological
 
 
+if TYPE_CHECKING:
+    from .dependency import DependencyProcessor
+    from .interfaces import MapperProperty
+    from .mapper import Mapper
+    from .session import Session
+    from .session import SessionTransaction
+    from .state import InstanceState
+
+
 def track_cascade_events(descriptor, prop):
     """Establish event listeners on object attributes which handle
     cascade-on-set/append.
@@ -131,7 +146,13 @@ def track_cascade_events(descriptor, prop):
 
 
 class UOWTransaction:
-    def __init__(self, session):
+    session: Session
+    transaction: SessionTransaction
+    attributes: Dict[str, Any]
+    deps: util.defaultdict[Mapper[Any], Set[DependencyProcessor]]
+    mappers: util.defaultdict[Mapper[Any], Set[InstanceState[Any]]]
+
+    def __init__(self, session: Session):
         self.session = session
 
         # dictionary used by external actors to
@@ -275,13 +296,13 @@ class UOWTransaction:
 
     def register_object(
         self,
-        state,
-        isdelete=False,
-        listonly=False,
-        cancel_delete=False,
-        operation=None,
-        prop=None,
-    ):
+        state: InstanceState[Any],
+        isdelete: bool = False,
+        listonly: bool = False,
+        cancel_delete: bool = False,
+        operation: Optional[str] = None,
+        prop: Optional[MapperProperty] = None,
+    ) -> bool:
         if not self.session._contains_state(state):
             # this condition is normal when objects are registered
             # as part of a relationship cascade operation.  it should
@@ -408,7 +429,7 @@ class UOWTransaction:
             [a for a in self.postsort_actions.values() if not a.disabled]
         ).difference(cycles)
 
-    def execute(self):
+    def execute(self) -> None:
         postsort_actions = self._generate_actions()
 
         postsort_actions = sorted(
@@ -435,7 +456,7 @@ class UOWTransaction:
             for rec in topological.sort(self.dependencies, postsort_actions):
                 rec.execute(self)
 
-    def finalize_flush_changes(self):
+    def finalize_flush_changes(self) -> None:
         """Mark processed objects as clean / deleted after a successful
         flush().