1 files changed, 60 insertions, 6 deletions

diff --git a/deluge/decorators.py b/deluge/decorators.py
index a601bf7e0..2f9fcd7d3 100644
--- a/deluge/decorators.py
+++ b/deluge/decorators.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 #
 # Copyright (C) 2010 John Garland <johnnybg+deluge@gmail.com>
 #
@@ -7,12 +6,13 @@
 # See LICENSE for more details.
 #
 
-from __future__ import unicode_literals
-
 import inspect
 import re
 import warnings
 from functools import wraps
+from typing import Any, Callable, Coroutine, TypeVar
+
+from twisted.internet import defer
 
 
 def proxy(proxy_func):
@@ -127,7 +127,7 @@ def _overrides(stack, method, explicit_base_classes=None):
                         % (
                             method.__name__,
                             cls,
-                            'File: %s:%s' % (stack[1][1], stack[1][2]),
+                            f'File: {stack[1][1]}:{stack[1][2]}',
                         )
                     )
 
@@ -137,7 +137,7 @@ def _overrides(stack, method, explicit_base_classes=None):
             % (
                 method.__name__,
                 check_classes,
-                'File: %s:%s' % (stack[1][1], stack[1][2]),
+                f'File: {stack[1][1]}:{stack[1][2]}',
             )
         )
     return method
@@ -154,7 +154,7 @@ def deprecated(func):
     def depr_func(*args, **kwargs):
         warnings.simplefilter('always', DeprecationWarning)  # Turn off filter
         warnings.warn(
-            'Call to deprecated function {}.'.format(func.__name__),
+            f'Call to deprecated function {func.__name__}.',
             category=DeprecationWarning,
             stacklevel=2,
         )
@@ -162,3 +162,57 @@ def deprecated(func):
         return func(*args, **kwargs)
 
     return depr_func
+
+
+class CoroutineDeferred(defer.Deferred):
+    """Wraps a coroutine in a Deferred.
+    It will dynamically pass through the underlying coroutine without wrapping where apporpriate."""
+
+    def __init__(self, coro: Coroutine):
+        # Delay this import to make sure a reactor was installed first
+        from twisted.internet import reactor
+
+        super().__init__()
+        self.coro = coro
+        self.awaited = None
+        self.activate_deferred = reactor.callLater(0, self.activate)
+
+    def __await__(self):
+        if self.awaited in [None, True]:
+            self.awaited = True
+            return self.coro.__await__()
+        # Already in deferred mode
+        return super().__await__()
+
+    def activate(self):
+        """If the result wasn't awaited before the next context switch, we turn it into a deferred."""
+        if self.awaited is None:
+            self.awaited = False
+            try:
+                d = defer.Deferred.fromCoroutine(self.coro)
+            except AttributeError:
+                # Fallback for Twisted <= 21.2 without fromCoroutine
+                d = defer.ensureDeferred(self.coro)
+            d.chainDeferred(self)
+
+    def addCallbacks(self, *args, **kwargs):  # noqa: N802
+        assert not self.awaited, 'Cannot add callbacks to an already awaited coroutine.'
+        self.activate()
+        return super().addCallbacks(*args, **kwargs)
+
+
+_RetT = TypeVar('_RetT')
+
+
+def maybe_coroutine(
+    f: Callable[..., Coroutine[Any, Any, _RetT]]
+) -> 'Callable[..., defer.Deferred[_RetT]]':
+    """Wraps a coroutine function to make it usable as a normal function that returns a Deferred."""
+
+    @wraps(f)
+    def wrapper(*args, **kwargs):
+        # Uncomment for quick testing to make sure CoroutineDeferred magic isn't at fault
+        # return defer.ensureDeferred(f(*args, **kwargs))
+        return CoroutineDeferred(f(*args, **kwargs))
+
+    return wrapper