ENH: Added netmask and hostmask (#30)

Author: TomAugspurger

cyberpandas/ip_array.py

@@ -535,6 +535,78 @@ def packed(self):
         # TODO: I wonder if that should be post-fixed by 0s.
         return self.data.tobytes()
 
+    def _apply_mask(self, op, v4_prefixlen, v6_prefixlen):
+        """Apply a netmask or hostmask"""
+        self = self.copy()
+        is_v4 = self.is_ipv4
+        v4_net = getattr(
+            ipaddress.ip_network(u'0.0.0.0/{}'.format(v4_prefixlen)),
+            op)
+        v4_mask = IPArray([v4_net])
+        self.data[is_v4] = v4_mask.data
+
+        v6_net = getattr(
+            ipaddress.ip_network(u'0::0/{}'.format(v6_prefixlen)),
+            op)
+        v6_mask = IPArray([v6_net])
+        self.data[~is_v4] = v6_mask.data
+        return self
+
+    def netmask(self, v4_prefixlen=32, v6_prefixlen=128):
+        """Compute an array of netmasks for an array of IP addresses.
+
+        Note that this is a method, rather than a property, to support
+        taking `v4_prefixlen` and `v6_prefixlen` as arguments.
+
+        Parameters
+        ----------
+        v4_prefixlen : int, default 32
+            Length of the network prefix, in bits, for IPv4 addresses
+        v6_prefixlen : int, default 128
+            Lnegth of the network prefix, in bits, for IPv6 addresses
+
+        Returns
+        -------
+        IPArray
+
+        See Also
+        --------
+        IPArray.hostmask
+
+        Examples
+        --------
+        >>> arr = ip.IPArray(['192.0.0.0', '1:1::'])
+        >>> arr.netmask(v4_prefixlen=16, v6_prefixlen=32)
+        IPArray(['255.255.0.0', 'ffff:ffff::'])
+        """
+        return self._apply_mask('netmask', v4_prefixlen, v6_prefixlen)
+
+    def hostmask(self, v4_prefixlen=32, v6_prefixlen=128):
+        """Compute an array of hostmasks for an array of IP addresses.
+
+        Parameters
+        ----------
+        v4_prefixlen : int, default 32
+            Length of the network prefix, in bits, for IPv4 addresses
+        v6_prefixlen : int, default 128
+            Lnegth of the network prefix, in bits, for IPv6 addresses
+
+        Returns
+        -------
+        IPArray
+
+        See Also
+        --------
+        IPArray.netmask
+
+        Examples
+        --------
+        >>> arr = ip.IPArray(['192.0.0.0', '1:1::'])
+        >>> arr.netmask(v4_prefixlen=16, v6_prefixlen=32)
+        IPArray(['0.0.255.255', '::ffff:ffff:ffff:ffff:ffff:ffff'])
+        """
+        return self._apply_mask('hostmask', v4_prefixlen, v6_prefixlen)
+
 
 # -----------------------------------------------------------------------------
 # Accessor
@@ -573,6 +645,14 @@ def isin(self, other):
         return delegated_method(self._data.isin, self._index,
                                 self._name, other)
 
+    def netmask(self, v4_prefixlen=32, v6_prefixlen=128):
+        return delegated_method(self._data.netmask, self._index,
+                                self._name, v4_prefixlen, v6_prefixlen)
+
+    def hostmask(self, v4_prefixlen=32, v6_prefixlen=128):
+        return delegated_method(self._data.hostmask, self._index,
+                                self._name, v4_prefixlen, v6_prefixlen)
+
 
 def is_ipaddress_type(obj):
     t = getattr(obj, 'dtype', obj)

docs/source/api.rst

@@ -39,11 +39,12 @@ an IPArray, calling the Series method will dispatch to these methods.
 .. automethod:: IPArray.unique
 .. automethod:: IPArray.isin
 .. automethod:: IPArray.isna
+.. autoattribute:: IPArray.netmask
 
 IP Address Attributes
 """""""""""""""""""""
 
-IP-addres specific attributes.
+IP addresss-specific attributes.
 
 .. autoattribute:: IPArray.is_ipv4
 .. autoattribute:: IPArray.is_ipv6
@@ -55,6 +56,9 @@ IP-addres specific attributes.
 .. autoattribute:: IPArray.is_reserved
 .. autoattribute:: IPArray.is_loopback
 .. autoattribute:: IPArray.is_link_local
+.. automethod:: IPArray.netmask
+.. automethod:: IPArray.hostmask
+
 
 
 :class:`MACArray`

tests/ip/test_ip.py

@@ -327,3 +327,58 @@ def test_ip_range(start, stop, step, expected):
     result = ip.ip_range(start, stop, step)
     expected = ip.IPArray(expected)
     assert result.equals(expected)
+
+
+@pytest.mark.parametrize('addresses', [
+    [u'0.0.0.0', u'192.168.1.1', u'::1:1:0:0:0:1']
+])
+@pytest.mark.parametrize('v4_prefixlen, v6_prefixlen', [
+    (32, 128),
+    (24, 96),
+    (16, 64),
+    (8, 8),
+])
+@pytest.mark.parametrize('op', ['netmask', 'hostmask'])
+def test_mask(op, v4_prefixlen, v6_prefixlen, addresses):
+    is_v6 = [':' in x for x in addresses]
+    prefixes = [v6_prefixlen if v6 else v4_prefixlen for v6 in is_v6]
+    networks = [
+        ipaddress.ip_network(u"{}/{}".format(addr, prefix), strict=False)
+        for addr, prefix in zip(addresses, prefixes)
+    ]
+    expected = [getattr(net, op) for net in networks]
+    call = operator.methodcaller(op, v4_prefixlen=v4_prefixlen,
+                                 v6_prefixlen=v6_prefixlen)
+    result = list(call(ip.IPArray(addresses)))
+    if op == 'hostmask':
+        # ipaddress will return an IPv6(0), which doesn't compare equal
+        # to an IPv4(0), our result.
+        expected = [int(x) for x in expected]
+        result = [int(x) for x in result]
+
+    assert result == expected
+
+
+def test_netmask_basic():
+    arr = ip.IPArray([u'192.0.0.0', u'1:1::'])
+    result = arr.netmask(v4_prefixlen=16, v6_prefixlen=32)
+    expected = ip.IPArray([u'255.255.0.0', u'ffff:ffff::'])
+    assert result.equals(expected)
+
+    result = pd.Series(arr, name='foo').ip.netmask(v4_prefixlen=16,
+                                                   v6_prefixlen=32)
+    assert result.name == 'foo'
+    assert result.values.equals(expected)
+
+
+def test_hostmask_basic():
+    arr = ip.IPArray([u'192.0.0.0', u'1:1::'])
+    result = arr.hostmask(v4_prefixlen=16, v6_prefixlen=32)
+    expected = ip.IPArray([u'0.0.255.255',
+                           u'::ffff:ffff:ffff:ffff:ffff:ffff'])
+    assert result.equals(expected)
+
+    result = pd.Series(arr, name='foo').ip.hostmask(v4_prefixlen=16,
+                                                    v6_prefixlen=32)
+    assert result.name == 'foo'
+    assert result.values.equals(expected)