From 32683690f4925e85b6b072d3b439919e861bf487 Mon Sep 17 00:00:00 2001
From: Steven Looman <steven.looman@gmail.com>
Date: Sun, 19 May 2024 22:34:32 +0200
Subject: [PATCH 1/2] Only fetch wanted IgdState items

---
 async_upnp_client/profiles/igd.py | 69 +++++++++++++++++++++++++++----
 1 file changed, 60 insertions(+), 9 deletions(-)

diff --git a/async_upnp_client/profiles/igd.py b/async_upnp_client/profiles/igd.py
index 2e892ff8..0224fb4a 100644
--- a/async_upnp_client/profiles/igd.py
+++ b/async_upnp_client/profiles/igd.py
@@ -4,8 +4,9 @@
 import asyncio
 import logging
 from datetime import datetime, timedelta
+from enum import Enum
 from ipaddress import IPv4Address
-from typing import List, NamedTuple, Optional, Sequence, Union, cast
+from typing import AbstractSet, List, NamedTuple, Optional, Sequence, Union, cast
 
 from async_upnp_client.client import UpnpAction, UpnpDevice
 from async_upnp_client.event_handler import UpnpEventHandler
@@ -93,7 +94,7 @@ class IgdState(NamedTuple):
     packets_received: Union[None, BaseException, int]
     packets_sent: Union[None, BaseException, int]
     status_info: Union[None, BaseException, StatusInfo]
-    external_ip_address: Union[str, BaseException, None]
+    external_ip_address: Union[None, BaseException, str]
 
     # Derived values.
     kibibytes_per_sec_received: Union[None, float]
@@ -102,6 +103,27 @@ class IgdState(NamedTuple):
     packets_per_sec_sent: Union[None, float]
 
 
+class IgdStateItem(Enum):
+    """
+    IGD state item.
+
+    Used to specify what to request from the device.
+    """
+
+    BYTES_RECEIVED = 1
+    BYTES_SENT = 2
+    PACKETS_RECEIVED = 3
+    PACKETS_SENT = 4
+
+    STATUS_INFO = 5
+    EXTERNAL_IP_ADDRESS = 6
+
+    KIBIBYTES_PER_SEC_RECEIVED = 11
+    KIBIBYTES_PER_SEC_SENT = 12
+    PACKETS_PER_SEC_RECEIVED = 13
+    PACKETS_PER_SEC_SENT = 14
+
+
 def _derive_value_per_second(
     value_name: str,
     current_timestamp: datetime,
@@ -614,7 +636,7 @@ async def async_set_default_connection_service(self, service: str) -> None:
         await action.async_call(NewDefaultConnectionService=service)
 
     async def async_get_traffic_and_status_data(
-        self,
+        self, items: Optional[AbstractSet[IgdStateItem]] = None
     ) -> IgdState:
         """
         Get all traffic data at once, including derived data.
@@ -632,14 +654,42 @@ async def async_get_traffic_and_status_data(
           * connection status
           * uptime
         """
+        items = items or {
+            IgdStateItem.BYTES_RECEIVED,
+            IgdStateItem.BYTES_SENT,
+            IgdStateItem.PACKETS_RECEIVED,
+            IgdStateItem.PACKETS_SENT,
+            IgdStateItem.STATUS_INFO,
+            IgdStateItem.EXTERNAL_IP_ADDRESS,
+        }
+
+        async def nop() -> None:
+            """Pass."""
+
         timestamp = datetime.now()
         values = await asyncio.gather(
-            self.async_get_total_bytes_received(),
-            self.async_get_total_bytes_sent(),
-            self.async_get_total_packets_received(),
-            self.async_get_total_packets_sent(),
-            self.async_get_status_info(),
-            self.async_get_external_ip_address(),
+            self.async_get_total_bytes_received()
+            if IgdStateItem.BYTES_RECEIVED in items
+            or IgdStateItem.KIBIBYTES_PER_SEC_RECEIVED in items
+            else nop(),
+            self.async_get_total_bytes_sent()
+            if IgdStateItem.BYTES_SENT in items
+            or IgdStateItem.KIBIBYTES_PER_SEC_SENT in items
+            else nop(),
+            self.async_get_total_packets_received()
+            if IgdStateItem.PACKETS_RECEIVED in items
+            or IgdStateItem.PACKETS_PER_SEC_RECEIVED in items
+            else nop(),
+            self.async_get_total_packets_sent()
+            if IgdStateItem.PACKETS_SENT in items
+            or IgdStateItem.PACKETS_PER_SEC_SENT in items
+            else nop(),
+            self.async_get_status_info()
+            if IgdStateItem.STATUS_INFO in items
+            else nop(),
+            self.async_get_external_ip_address()
+            if IgdStateItem.EXTERNAL_IP_ADDRESS in items
+            else nop(),
             return_exceptions=True,
         )
 
@@ -684,6 +734,7 @@ async def async_get_traffic_and_status_data(
             packets_sent_original=values[3],
         )
 
+        # Test if any of the calls were ok. If not, raise the exception.
         non_exceptions = [
             value for value in values if not isinstance(value, BaseException)
         ]

From 7e30a2c0c5b82c773f071d29c2c132e2e0aa694b Mon Sep 17 00:00:00 2001
From: Steven Looman <steven.looman@gmail.com>
Date: Sun, 19 May 2024 22:35:53 +0200
Subject: [PATCH 2/2] Add change item

---
 changes/227.feature | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changes/227.feature

diff --git a/changes/227.feature b/changes/227.feature
new file mode 100644
index 00000000..61a39399
--- /dev/null
+++ b/changes/227.feature
@@ -0,0 +1 @@
+Only fetch wanted IgdState items in IGD profile.