Appease mypy

1 files changed, 18 insertions, 13 deletions

diff --git a/synapse/util/patch_inline_callbacks.py b/synapse/util/patch_inline_callbacks.py
index b518dae256..64a2c891c3 100644
--- a/synapse/util/patch_inline_callbacks.py
+++ b/synapse/util/patch_inline_callbacks.py
@@ -18,11 +18,17 @@ from __future__ import print_function
 import functools
 import sys
 
+from typing import List, Callable, Any
+
 from twisted.internet import defer
 from twisted.internet.defer import Deferred
 from twisted.python.failure import Failure
 
 
+# Tracks if we've already patched inlineCallbacks
+_already_patched = False
+
+
 def do_patch():
     """
     Patch defer.inlineCallbacks so that it checks the state of the logcontext on exit
@@ -30,16 +36,18 @@ def do_patch():
 
     from synapse.logging.context import LoggingContext
 
+    global _already_patched
+
     orig_inline_callbacks = defer.inlineCallbacks
-    if hasattr(orig_inline_callbacks, "patched_by_synapse"):
+    if _already_patched:
         return
 
     def new_inline_callbacks(f):
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
             start_context = LoggingContext.current_context()
-            changes = []
-            orig = orig_inline_callbacks(_check_yield_points(f, changes, start_context))
+            changes: List[str] = []
+            orig = orig_inline_callbacks(_check_yield_points(f, changes))
 
             try:
                 res = orig(*args, **kwargs)
@@ -101,10 +109,10 @@ def do_patch():
         return wrapped
 
     defer.inlineCallbacks = new_inline_callbacks
-    new_inline_callbacks.patched_by_synapse = True
+    _already_patched = True
 
 
-def _check_yield_points(f, changes, start_context):
+def _check_yield_points(f: Callable, changes: List[str]):
     """Wraps a generator that is about to be passed to defer.inlineCallbacks
     checking that after every yield the log contexts are correct.
 
@@ -114,9 +122,8 @@ def _check_yield_points(f, changes, start_context):
 
     Args:
         f: generator function to wrap
-        changes (list[str]): A list of strings detailing how the contexts
+        changes: A list of strings detailing how the contexts
             changed within a function.
-        start_context (LoggingContext): The initial context we're expecting
 
     Returns:
         function
@@ -126,13 +133,13 @@ def _check_yield_points(f, changes, start_context):
 
     @functools.wraps(f)
     def check_yield_points_inner(*args, **kwargs):
-        expected_context = start_context
-
         gen = f(*args, **kwargs)
 
         last_yield_line_no = gen.gi_frame.f_lineno
-        result = None
+        result: Any = None
         while True:
+            expected_context = LoggingContext.current_context()
+
             try:
                 isFailure = isinstance(result, Failure)
                 if isFailure:
@@ -200,7 +207,7 @@ def _check_yield_points(f, changes, start_context):
                     "%s changed context from %s to %s, happened between lines %d and %d in %s"
                     % (
                         frame.f_code.co_name,
-                        start_context,
+                        expected_context,
                         LoggingContext.current_context(),
                         last_yield_line_no,
                         frame.f_lineno,
@@ -209,8 +216,6 @@ def _check_yield_points(f, changes, start_context):
                 )
                 changes.append(err)
 
-                expected_context = LoggingContext.current_context()
-
             last_yield_line_no = frame.f_lineno
 
     return check_yield_points_inner