feat: improve performance of ip address caching (#1392)

bdraco · web-flow · commit f7c77081b2f8 · 2024-08-26T15:50:02.000-10:00
diff --git a/src/zeroconf/_utils/ipaddress.py b/src/zeroconf/_utils/ipaddress.py
@@ -28,20 +28,33 @@
 from .._dns import DNSAddress
 from ..const import _TYPE_AAAA
 
+if sys.version_info >= (3, 9, 0):
+    from functools import cache
+else:
+    cache = lru_cache(maxsize=None)
+
 bytes_ = bytes
 int_ = int
 IPADDRESS_SUPPORTS_SCOPE_ID = sys.version_info >= (3, 9, 0)
 
 
 class ZeroconfIPv4Address(IPv4Address):
-    __slots__ = ("_str", "_is_link_local", "_is_unspecified")
+    __slots__ = (
+        "_str",
+        "_is_link_local",
+        "_is_unspecified",
+        "_is_loopback",
+        "__hash__",
+    )
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize a new IPv4 address."""
         super().__init__(*args, **kwargs)
         self._str = super().__str__()
         self._is_link_local = super().is_link_local
         self._is_unspecified = super().is_unspecified
+        self._is_loopback = super().is_loopback
+        self.__hash__ = cache(lambda: IPv4Address.__hash__(self))  # type: ignore[method-assign]
 
     def __str__(self) -> str:
         """Return the string representation of the IPv4 address."""
@@ -57,16 +70,29 @@ def is_unspecified(self) -> bool:
         """Return True if this is an unspecified address."""
         return self._is_unspecified
 
+    @property
+    def is_loopback(self) -> bool:
+        """Return True if this is a loop back."""
+        return self._is_loopback
+
 
 class ZeroconfIPv6Address(IPv6Address):
-    __slots__ = ("_str", "_is_link_local", "_is_unspecified")
+    __slots__ = (
+        "_str",
+        "_is_link_local",
+        "_is_unspecified",
+        "_is_loopback",
+        "__hash__",
+    )
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         """Initialize a new IPv6 address."""
         super().__init__(*args, **kwargs)
         self._str = super().__str__()
         self._is_link_local = super().is_link_local
         self._is_unspecified = super().is_unspecified
+        self._is_loopback = super().is_loopback
+        self.__hash__ = cache(lambda: IPv6Address.__hash__(self))  # type: ignore[method-assign]
 
     def __str__(self) -> str:
         """Return the string representation of the IPv6 address."""
@@ -82,6 +108,11 @@ def is_unspecified(self) -> bool:
         """Return True if this is an unspecified address."""
         return self._is_unspecified
 
+    @property
+    def is_loopback(self) -> bool:
+        """Return True if this is a loop back."""
+        return self._is_loopback
+
 
 @lru_cache(maxsize=512)
 def _cached_ip_addresses(
diff --git a/tests/utils/test_ipaddress.py b/tests/utils/test_ipaddress.py
@@ -19,7 +19,17 @@ def test_cached_ip_addresses_wrapper():
         str(ipaddress.cached_ip_addresses(b"&\x06(\x00\x02 \x00\x01\x02H\x18\x93%\xc8\x19F"))
         == "2606:2800:220:1:248:1893:25c8:1946"
     )
-    assert ipaddress.cached_ip_addresses("::1") == ipaddress.IPv6Address("::1")
+    loop_back_ipv6 = ipaddress.cached_ip_addresses("::1")
+    assert loop_back_ipv6 == ipaddress.IPv6Address("::1")
+    assert loop_back_ipv6.is_loopback is True
+
+    assert hash(loop_back_ipv6) == hash(ipaddress.IPv6Address("::1"))
+
+    loop_back_ipv4 = ipaddress.cached_ip_addresses("127.0.0.1")
+    assert loop_back_ipv4 == ipaddress.IPv4Address("127.0.0.1")
+    assert loop_back_ipv4.is_loopback is True
+
+    assert hash(loop_back_ipv4) == hash(ipaddress.IPv4Address("127.0.0.1"))
 
     ipv4 = ipaddress.cached_ip_addresses("169.254.0.0")
     assert ipv4 is not None
