From 3e2094b59f90465cb6c558639b2d348e209c75b1 Mon Sep 17 00:00:00 2001
From: mdeweerd <mdeweerd@users.noreply.github.com>
Date: Thu, 6 Jul 2023 16:59:54 +0200
Subject: [PATCH] Update for zigpy >= 0.56.0, HA 2023.7.0 - integrate retryable

---
 custom_components/zha_toolkit/neighbours.py  | 16 +-----
 custom_components/zha_toolkit/ota.py         |  8 +--
 custom_components/zha_toolkit/scan_device.py |  5 +-
 custom_components/zha_toolkit/utils.py       | 59 +++++++++++++++++++-
 4 files changed, 65 insertions(+), 23 deletions(-)

diff --git a/custom_components/zha_toolkit/neighbours.py b/custom_components/zha_toolkit/neighbours.py
index cb62bf0..9e03dbd 100644
--- a/custom_components/zha_toolkit/neighbours.py
+++ b/custom_components/zha_toolkit/neighbours.py
@@ -8,25 +8,11 @@
 
 import zigpy.zdo.types as zdo_t
 from homeassistant.util.json import save_json
-from zigpy.exceptions import ControllerException, DeliveryError
-from zigpy.util import retryable
+from zigpy.exceptions import DeliveryError
 
 LOGGER = logging.getLogger(__name__)
 
 
-@retryable(
-    (
-        DeliveryError,
-        ControllerException,
-        asyncio.CancelledError,
-        asyncio.TimeoutError,
-    ),
-    tries=5,
-)
-def wrapper(cmd, *args, **kwargs):
-    return cmd(*args, **kwargs)
-
-
 async def routes_and_neighbours(
     app, listener, ieee, cmd, data, service, params, event_data
 ):
diff --git a/custom_components/zha_toolkit/ota.py b/custom_components/zha_toolkit/ota.py
index 74ab13f..15aae22 100644
--- a/custom_components/zha_toolkit/ota.py
+++ b/custom_components/zha_toolkit/ota.py
@@ -8,9 +8,9 @@
 from pkg_resources import parse_version
 from zigpy import __version__ as zigpy_version
 from zigpy.exceptions import ControllerException, DeliveryError
-from zigpy.util import retryable
 
 from . import DEFAULT_OTAU
+from . import utils as u
 from .params import INTERNAL_PARAMS as p
 
 LOGGER = logging.getLogger(__name__)
@@ -21,7 +21,7 @@
 SONOFF_LIST_URL = "https://zigbee-ota.sonoff.tech/releases/upgrade.json"
 
 
-@retryable(
+@u.retryable(
     (
         DeliveryError,
         ControllerException,
@@ -30,7 +30,7 @@
     ),
     tries=3,
 )
-async def wrapper(cmd, *args, **kwargs):
+async def retry_wrapper(cmd, *args, **kwargs):
     return await cmd(*args, **kwargs)
 
 
@@ -232,7 +232,7 @@ async def ota_notify(
         ret = await cluster.image_notify(0, 100)
     else:
         cmd_args = [0, 100]
-        ret = await wrapper(
+        ret = await retry_wrapper(
             cluster.client_command,
             0,  # cmd_id
             *cmd_args,
diff --git a/custom_components/zha_toolkit/scan_device.py b/custom_components/zha_toolkit/scan_device.py
index 5a433e6..f770ac5 100644
--- a/custom_components/zha_toolkit/scan_device.py
+++ b/custom_components/zha_toolkit/scan_device.py
@@ -6,7 +6,6 @@
 
 from zigpy import types as t
 from zigpy.exceptions import ControllerException, DeliveryError
-from zigpy.util import retryable
 from zigpy.zcl import foundation
 
 from . import utils as u
@@ -15,7 +14,7 @@
 LOGGER = logging.getLogger(__name__)
 
 
-@retryable(
+@u.retryable(
     (
         DeliveryError,
         ControllerException,
@@ -30,7 +29,7 @@ async def read_attr(cluster, attrs, manufacturer=None):
     )
 
 
-@retryable(
+@u.retryable(
     (DeliveryError, asyncio.CancelledError, asyncio.TimeoutError), tries=3
 )
 async def wrapper(cmd, *args, **kwargs):
diff --git a/custom_components/zha_toolkit/utils.py b/custom_components/zha_toolkit/utils.py
index b7bd792..77143b3 100644
--- a/custom_components/zha_toolkit/utils.py
+++ b/custom_components/zha_toolkit/utils.py
@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 import json
 import logging
 import os
 import re
 import typing
 from enum import Enum
+from typing import Any
 
 import packaging
 import packaging.version
@@ -15,7 +17,6 @@
 from zigpy import __version__ as zigpy_version
 from zigpy import types as t
 from zigpy.exceptions import ControllerException, DeliveryError
-from zigpy.util import retryable
 from zigpy.zcl import foundation as f
 
 from .params import INTERNAL_PARAMS as p
@@ -834,6 +835,62 @@ def extractParams(  # noqa: C901
     return params
 
 
+#
+# Copied retry and retryable from zigpy < 0.56.0
+#   where "tries" and "delay" were removed
+#  from the wrapper function and hence propagated to the decorated function.
+#
+async def retry(
+    func: typing.Callable[[], typing.Awaitable[typing.Any]],
+    retry_exceptions: typing.Iterable[Any],  # typing.Iterable[BaseException],
+    tries: int = 3,
+    delay: int | float = 0.1,
+) -> typing.Any:
+    """Retry a function in case of exception
+
+    Only exceptions in `retry_exceptions` will be retried.
+    """
+    while True:
+        LOGGER.debug("Tries remaining: %s", tries)
+        try:
+            return await func()
+        except retry_exceptions:  # type:ignore[misc]
+            if tries <= 1:
+                raise
+            tries -= 1
+            await asyncio.sleep(delay)
+
+
+def retryable(
+    retry_exceptions: typing.Iterable[Any],  # typing.Iterable[BaseException]
+    tries: int = 1,
+    delay: float = 0.1,
+) -> typing.Callable:
+    """Return a decorator which makes a function able to be retried
+
+    This adds "tries" and "delay" keyword arguments to the function. Only
+    exceptions in `retry_exceptions` will be retried.
+    """
+
+    def decorator(func: typing.Callable) -> typing.Callable:
+        nonlocal tries, delay
+
+        @functools.wraps(func)
+        def wrapper(*args, tries=tries, delay=delay, **kwargs):
+            if tries <= 1:
+                return func(*args, **kwargs)
+            return retry(
+                functools.partial(func, *args, **kwargs),
+                retry_exceptions,
+                tries=tries,
+                delay=delay,
+            )
+
+        return wrapper
+
+    return decorator
+
+
 # zigpy wrappers
 
 # The zigpy library does not offer retryable on read_attributes.