Linting

Author: andystaples

durabletask/entities/entity_lock.py

@@ -1,9 +1,17 @@
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from durabletask.task import OrchestrationContext
+
+
 class EntityLock:
-    def __init__(self, context):
+    # Note: This should
+    def __init__(self, context: 'OrchestrationContext'):
         self._context = context
 
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):  # TODO: Handle exceptions?
+    def __exit__(self, exc_type, exc_val, exc_tb):
         self._context._exit_critical_section()

durabletask/task.py

@@ -258,6 +258,10 @@ def continue_as_new(self, new_input: Any, *, save_events: bool = False) -> None:
         """
         pass
 
+    @abstractmethod
+    def _exit_critical_section(self) -> None:
+        pass
+
 
 class FailureDetails:
     def __init__(self, message: str, error_type: str, stack_trace: Optional[str]):

durabletask/worker.py

@@ -1169,7 +1169,7 @@ def lock_entities_function_helper(self, id: int, entities: list[EntityInstanceId
     def _exit_critical_section(self) -> None:
         if not self._entity_context.is_inside_critical_section:
             # Possible if the user calls continue_as_new inside the lock - in the success case, we will call
-            # _exit_critical_section both from the EntityLock and the exit logic. We must keep both calls in
+            # _exit_critical_section both from the EntityLock and the continue_as_new logic. We must keep both calls in
             # case the user code crashes after calling continue_as_new but before the EntityLock object is exited.
             return
         for entity_unlock_message in self._entity_context.emit_lock_release_messages():