Issue #16531: ipaddress.IPv4Network and ipaddress.IPv6Network now accept an (address, netmask) tuple argument, so as to easily construct network objects from existing addresses.

2 files changed, 187 insertions, 21 deletions

diff --git a/Lib/ipaddress.py b/Lib/ipaddress.py
index 1e451f8..216ddf1 100644
--- a/Lib/ipaddress.py
+++ b/Lib/ipaddress.py
@@ -991,15 +991,15 @@ class _BaseNetwork(_IPAddressBase):
                 raise ValueError('cannot set prefixlen_diff and new_prefix')
             prefixlen_diff = self._prefixlen - new_prefix
 
-        if self.prefixlen - prefixlen_diff < 0:
+        new_prefixlen = self.prefixlen - prefixlen_diff
+        if new_prefixlen < 0:
             raise ValueError(
                 'current prefixlen is %d, cannot have a prefixlen_diff of %d' %
                 (self.prefixlen, prefixlen_diff))
-        # TODO (pmoody): optimize this.
-        t = self.__class__('%s/%d' % (self.network_address,
-                                      self.prefixlen - prefixlen_diff),
-                                     strict=False)
-        return t.__class__('%s/%d' % (t.network_address, t.prefixlen))
+        return self.__class__((
+            int(self.network_address) & (int(self.netmask) << prefixlen_diff),
+            new_prefixlen
+            ))
 
     @property
     def is_multicast(self):
@@ -1389,6 +1389,18 @@ class IPv4Interface(IPv4Address):
             self._prefixlen = self._max_prefixlen
             return
 
+        if isinstance(address, tuple):
+            IPv4Address.__init__(self, address[0])
+            if len(address) > 1:
+                self._prefixlen = int(address[1])
+            else:
+                self._prefixlen = self._max_prefixlen
+
+            self.network = IPv4Network(address, strict=False)
+            self.netmask = self.network.netmask
+            self.hostmask = self.network.hostmask
+            return
+
         addr = _split_optional_netmask(address)
         IPv4Address.__init__(self, addr[0])
 
@@ -1504,22 +1516,42 @@ class IPv4Network(_BaseV4, _BaseNetwork):
         _BaseV4.__init__(self, address)
         _BaseNetwork.__init__(self, address)
 
-        # Constructing from a packed address
-        if isinstance(address, bytes):
+        # Constructing from a packed address or integer
+        if isinstance(address, (int, bytes)):
             self.network_address = IPv4Address(address)
             self._prefixlen = self._max_prefixlen
             self.netmask = IPv4Address(self._ALL_ONES)
             #fixme: address/network test here
             return
 
-        # Efficient constructor from integer.
-        if isinstance(address, int):
-            self.network_address = IPv4Address(address)
-            self._prefixlen = self._max_prefixlen
-            self.netmask = IPv4Address(self._ALL_ONES)
-            #fixme: address/network test here.
+        if isinstance(address, tuple):
+            if len(address) > 1:
+                # If address[1] is a string, treat it like a netmask.
+                if isinstance(address[1], str):
+                    self.netmask = IPv4Address(address[1])
+                    self._prefixlen = self._prefix_from_ip_int(
+                        int(self.netmask))
+                # address[1] should be an int.
+                else:
+                    self._prefixlen = int(address[1])
+                    self.netmask = IPv4Address(self._ip_int_from_prefix(
+                            self._prefixlen))
+            # We weren't given an address[1].
+            else:
+                self._prefixlen = self._max_prefixlen
+                self.netmask = IPv4Address(self._ip_int_from_prefix(
+                        self._prefixlen))
+            self.network_address = IPv4Address(address[0])
+            packed = int(self.network_address)
+            if packed & int(self.netmask) != packed:
+                if strict:
+                    raise ValueError('%s has host bits set' % self)
+                else:
+                    self.network_address = IPv4Address(packed &
+                                                       int(self.netmask))
             return
 
+
         # Assume input argument to be string or any object representation
         # which converts into a formatted IP prefix string.
         addr = _split_optional_netmask(address)
@@ -2030,6 +2062,16 @@ class IPv6Interface(IPv6Address):
             self.network = IPv6Network(self._ip)
             self._prefixlen = self._max_prefixlen
             return
+        if isinstance(address, tuple):
+            IPv6Address.__init__(self, address[0])
+            if len(address) > 1:
+                self._prefixlen = int(address[1])
+            else:
+                self._prefixlen = self._max_prefixlen
+            self.network = IPv6Network(address, strict=False)
+            self.netmask = self.network.netmask
+            self.hostmask = self.network.hostmask
+            return
 
         addr = _split_optional_netmask(address)
         IPv6Address.__init__(self, addr[0])
@@ -2147,18 +2189,29 @@ class IPv6Network(_BaseV6, _BaseNetwork):
         _BaseV6.__init__(self, address)
         _BaseNetwork.__init__(self, address)
 
-        # Efficient constructor from integer.
-        if isinstance(address, int):
+        # Efficient constructor from integer or packed address
+        if isinstance(address, (bytes, int)):
             self.network_address = IPv6Address(address)
             self._prefixlen = self._max_prefixlen
             self.netmask = IPv6Address(self._ALL_ONES)
             return
 
-        # Constructing from a packed address
-        if isinstance(address, bytes):
-            self.network_address = IPv6Address(address)
-            self._prefixlen = self._max_prefixlen
-            self.netmask = IPv6Address(self._ALL_ONES)
+        if isinstance(address, tuple):
+            self.network_address = IPv6Address(address[0])
+            if len(address) > 1:
+                self._prefixlen = int(address[1])
+            else:
+                self._prefixlen = self._max_prefixlen
+            self.netmask = IPv6Address(self._ip_int_from_prefix(
+                    self._prefixlen))
+            self.network_address = IPv6Address(address[0])
+            packed = int(self.network_address)
+            if packed & int(self.netmask) != packed:
+                if strict:
+                    raise ValueError('%s has host bits set' % self)
+                else:
+                    self.network_address = IPv6Address(packed &
+                                                       int(self.netmask))
             return
 
         # Assume input argument to be string or any object representation
diff --git a/Lib/test/test_ipaddress.py b/Lib/test/test_ipaddress.py
index 33adb3b..0b71bf8 100644
--- a/Lib/test/test_ipaddress.py
+++ b/Lib/test/test_ipaddress.py
@@ -628,6 +628,119 @@ class IpaddrUnitTest(unittest.TestCase):
         self.assertEqual("IPv6Interface('::1/128')",
                          repr(ipaddress.IPv6Interface('::1')))
 
+    # issue #16531: constructing IPv4Network from a (address, mask) tuple
+    def testIPv4Tuple(self):
+        # /32
+        ip = ipaddress.IPv4Address('192.0.2.1')
+        net = ipaddress.IPv4Network('192.0.2.1/32')
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.1', 32)), net)
+        self.assertEqual(ipaddress.IPv4Network((ip, 32)), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225985, 32)), net)
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.1',
+                                                '255.255.255.255')), net)
+        self.assertEqual(ipaddress.IPv4Network((ip,
+                                                '255.255.255.255')), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225985,
+                                                '255.255.255.255')), net)
+        # strict=True and host bits set
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network(('192.0.2.1', 24))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network((ip, 24))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network((3221225985, 24))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network(('192.0.2.1', '255.255.255.0'))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network((ip, '255.255.255.0'))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv4Network((3221225985, '255.255.255.0'))
+        # strict=False and host bits set
+        net = ipaddress.IPv4Network('192.0.2.0/24')
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.1', 24),
+                                               strict=False), net)
+        self.assertEqual(ipaddress.IPv4Network((ip, 24),
+                                               strict=False), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225985, 24),
+                                               strict=False), net)
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.1',
+                                                '255.255.255.0'),
+                                               strict=False), net)
+        self.assertEqual(ipaddress.IPv4Network((ip,
+                                                '255.255.255.0'),
+                                               strict=False), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225985,
+                                                '255.255.255.0'),
+                                               strict=False), net)
+
+        # /24
+        ip = ipaddress.IPv4Address('192.0.2.0')
+        net = ipaddress.IPv4Network('192.0.2.0/24')
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.0',
+                                                '255.255.255.0')), net)
+        self.assertEqual(ipaddress.IPv4Network((ip,
+                                                '255.255.255.0')), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225984,
+                                                '255.255.255.0')), net)
+        self.assertEqual(ipaddress.IPv4Network(('192.0.2.0', 24)), net)
+        self.assertEqual(ipaddress.IPv4Network((ip, 24)), net)
+        self.assertEqual(ipaddress.IPv4Network((3221225984, 24)), net)
+
+        self.assertEqual(ipaddress.IPv4Interface(('192.0.2.1', 24)),
+                         ipaddress.IPv4Interface('192.0.2.1/24'))
+        self.assertEqual(ipaddress.IPv4Interface((3221225985, 24)),
+                         ipaddress.IPv4Interface('192.0.2.1/24'))
+
+    # issue #16531: constructing IPv6Network from a (address, mask) tuple
+    def testIPv6Tuple(self):
+        # /128
+        ip = ipaddress.IPv6Address('2001:db8::')
+        net = ipaddress.IPv6Network('2001:db8::/128')
+        self.assertEqual(ipaddress.IPv6Network(('2001:db8::', '128')),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network(
+                (42540766411282592856903984951653826560, 128)),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network((ip, '128')),
+                         net)
+        ip = ipaddress.IPv6Address('2001:db8::')
+        net = ipaddress.IPv6Network('2001:db8::/96')
+        self.assertEqual(ipaddress.IPv6Network(('2001:db8::', '96')),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network(
+                (42540766411282592856903984951653826560, 96)),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network((ip, '96')),
+                         net)
+
+        # strict=True and host bits set
+        ip = ipaddress.IPv6Address('2001:db8::1')
+        with self.assertRaises(ValueError):
+            ipaddress.IPv6Network(('2001:db8::1', 96))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv6Network((
+                42540766411282592856903984951653826561, 96))
+        with self.assertRaises(ValueError):
+            ipaddress.IPv6Network((ip, 96))
+        # strict=False and host bits set
+        net = ipaddress.IPv6Network('2001:db8::/96')
+        self.assertEqual(ipaddress.IPv6Network(('2001:db8::1', 96),
+                                               strict=False),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network(
+                             (42540766411282592856903984951653826561, 96),
+                             strict=False),
+                         net)
+        self.assertEqual(ipaddress.IPv6Network((ip, 96), strict=False),
+                         net)
+
+        # /96
+        self.assertEqual(ipaddress.IPv6Interface(('2001:db8::1', '96')),
+                         ipaddress.IPv6Interface('2001:db8::1/96'))
+        self.assertEqual(ipaddress.IPv6Interface(
+                (42540766411282592856903984951653826561, '96')),
+                         ipaddress.IPv6Interface('2001:db8::1/96'))
+
     # issue57
     def testAddressIntMath(self):
         self.assertEqual(ipaddress.IPv4Address('1.1.1.1') + 255,