feat: add ip_addresses_by_version to ServiceInfo (#1145)

bdraco · web-flow · commit 524494edd49b · 2023-04-01T20:05:37.000-10:00
diff --git a/src/zeroconf/_services/info.py b/src/zeroconf/_services/info.py
@@ -233,6 +233,16 @@ def addresses_by_version(self, version: IPVersion) -> List[bytes]:
             *(addr.packed for addr in self._ipv6_addresses),
         ]
 
+    def ip_addresses_by_version(
+        self, version: IPVersion
+    ) -> Union[List[ipaddress.IPv4Address], List[ipaddress.IPv6Address], List[ipaddress._BaseAddress]]:
+        """List ip_address objects matching IP version."""
+        if version == IPVersion.V4Only:
+            return self._ipv4_addresses
+        if version == IPVersion.V6Only:
+            return self._ipv6_addresses
+        return [*self._ipv4_addresses, *self._ipv6_addresses]
+
     def parsed_addresses(self, version: IPVersion = IPVersion.All) -> List[str]:
         """List addresses in their parsed string form."""
         result = self.addresses_by_version(version)
diff --git a/tests/services/test_info.py b/tests/services/test_info.py
@@ -9,6 +9,7 @@
 import socket
 import threading
 import unittest
+from ipaddress import ip_address
 from threading import Event
 from typing import Iterable, List, Optional
 from unittest.mock import patch
@@ -540,8 +541,18 @@ def test_multiple_addresses():
         for info in infos:
             assert info.addresses == [address]
             assert info.addresses_by_version(r.IPVersion.All) == [address, address_v6, address_v6_ll]
+            assert info.ip_addresses_by_version(r.IPVersion.All) == [
+                ip_address(address),
+                ip_address(address_v6),
+                ip_address(address_v6_ll),
+            ]
             assert info.addresses_by_version(r.IPVersion.V4Only) == [address]
+            assert info.ip_addresses_by_version(r.IPVersion.V4Only) == [ip_address(address)]
             assert info.addresses_by_version(r.IPVersion.V6Only) == [address_v6, address_v6_ll]
+            assert info.ip_addresses_by_version(r.IPVersion.V6Only) == [
+                ip_address(address_v6),
+                ip_address(address_v6_ll),
+            ]
             assert info.parsed_addresses() == [address_parsed, address_v6_parsed, address_v6_ll_parsed]
             assert info.parsed_addresses(r.IPVersion.V4Only) == [address_parsed]
             assert info.parsed_addresses(r.IPVersion.V6Only) == [address_v6_parsed, address_v6_ll_parsed]
