refactor: move decorators under MultiProvider class

lidofinance · May 10, 2024 · 7a451f0 · 7a451f0
1 parent 80f7d9e
commit 7a451f0
Showing 1 changed file with 37 additions and 35 deletions.
diff --git a/src/providers/ipfs/multi.py b/src/providers/ipfs/multi.py
@@ -34,41 +34,6 @@ class MaxRetryError(IPFSError):
     ...
 
 
-def retry(fn):
-    @wraps(fn)
-    def wrapped(self: SupportsRetries, *args, **kwargs):
-        retries_left = self.retries
-        while retries_left:
-            try:
-                return fn(self, *args, **kwargs)
-            except IPFSError as ex:
-                retries_left -= 1
-                if not retries_left:
-                    raise MaxRetryError from ex
-                logger.warning({"msg": f"Retrying a failed call of {fn.__name__}, {retries_left=}", "error": str(ex)})
-        raise MaxRetryError
-
-    return wrapped
-
-
-def with_fallback(fn):
-    @wraps(fn)
-    def wrapped(self: MultiProvider, *args, **kwargs):
-        try:
-            result = fn(self, *args, **kwargs)
-        except IPFSError:
-            self.current_provider_index = (self.current_provider_index + 1) % len(self.providers)
-            if self.last_working_provider_index == self.current_provider_index:
-                logger.error({"msg": "No more IPFS providers left to call"})
-                raise
-            return wrapped(self, *args, **kwargs)
-
-        self.last_working_provider_index = self.current_provider_index
-        return result
-
-    return wrapped
-
-
 class MultiIPFSProvider(IPFSProvider, MultiProvider[IPFSProvider]):
     """Fallback-driven provider for IPFS"""
 
@@ -82,6 +47,43 @@ def __init__(self, providers: Iterable[IPFSProvider], *, retries: int = 3) -> No
         for p in self.providers:
             assert isinstance(p, IPFSProvider)
 
+    @staticmethod
+    def with_fallback(fn):
+        @wraps(fn)
+        def wrapped(self: MultiProvider, *args, **kwargs):
+            try:
+                result = fn(self, *args, **kwargs)
+            except IPFSError:
+                self.current_provider_index = (self.current_provider_index + 1) % len(self.providers)
+                if self.last_working_provider_index == self.current_provider_index:
+                    logger.error({"msg": "No more IPFS providers left to call"})
+                    raise
+                return wrapped(self, *args, **kwargs)
+
+            self.last_working_provider_index = self.current_provider_index
+            return result
+
+        return wrapped
+
+    @staticmethod
+    def retry(fn):
+        @wraps(fn)
+        def wrapped(self: SupportsRetries, *args, **kwargs):
+            retries_left = self.retries
+            while retries_left:
+                try:
+                    return fn(self, *args, **kwargs)
+                except IPFSError as ex:
+                    retries_left -= 1
+                    if not retries_left:
+                        raise MaxRetryError from ex
+                    logger.warning(
+                        {"msg": f"Retrying a failed call of {fn.__name__}, {retries_left=}", "error": str(ex)}
+                    )
+            raise MaxRetryError
+
+        return wrapped
+
     @with_fallback
     @retry
     def fetch(self, cid: CIDv0 | CIDv1) -> bytes: