ENH: Added ip_range (#27)

Author: TomAugspurger

cyberpandas/__init__.py

@@ -5,6 +5,7 @@
     IPArray,
     IPAccessor,
 )
+from .ip_methods import ip_range
 from .parser import to_ipaddress
 from .mac_array import MACType, MACArray
 
@@ -26,5 +27,6 @@
     'IPType',
     'MACArray',
     'MACType',
+    'ip_range',
     'to_ipaddress',
 ]

cyberpandas/ip_methods.py

@@ -0,0 +1,70 @@
+import ipaddress
+
+import numpy as np
+import six
+
+from .ip_array import IPArray
+from .common import _U8_MAX
+
+
+def _as_int(ip):
+    # type: (Union[int, str, IPv4Address, IPv6Address]) -> int
+    if isinstance(ip, six.string_types):
+        ip = ipaddress.ip_address(ip)
+    return int(ip)
+
+
+def _crosses_boundary(lo, hi):
+    return (lo <= _U8_MAX) == (hi <= _U8_MAX)
+
+
+def ip_range(start=None, stop=None, step=None):
+    """Generate a range of IP Addresses
+
+    Parameters
+    ----------
+    start : int, str, IPv4Address, or IPv6Address, optional
+        Start of interval.  The interval includes this value.  The default
+        start value is 0.
+    start : int, str, IPv4Address, or IPv6Address, optional
+        End of interval.  The interval does not include this value.
+    step : int, optional
+        Spacing between values.  For any output `out`, this is the distance
+        between two adjacent values, ``out[i+1] - out[i]``.  The default
+        step size is 1.  If `step` is specified as a position argument,
+        `start` must also be given.
+
+    Returns
+    -------
+    IPArray
+
+    Notes
+    -----
+    Performance will worsen if either of `start` or `stop` are larger than
+    2**64.
+
+    Examples
+    --------
+    From integers
+
+    >>> ip_range(1, 5)
+    IPArray(['0.0.0.1', '0.0.0.2', '0.0.0.3', '0.0.0.4'])
+
+    Or strings
+
+    >>> ip_range('0.0.0.1', '0.0.0.5')
+    IPArray(['0.0.0.1', '0.0.0.2', '0.0.0.3', '0.0.0.4'])
+
+    Or `ipaddress` objects
+
+    >>> ip_range(ipaddress.IPv4Address(1), ipaddress.IPv4Address(5))
+    IPArray(['0.0.0.1', '0.0.0.2', '0.0.0.3', '0.0.0.4'])
+    """
+    if start is not None:
+        start = _as_int(start)
+    if stop is not None:
+        stop = _as_int(stop)
+    if step is not None:
+        step = _as_int(step)
+    arr = IPArray(np.arange(start, stop, step))
+    return arr

cyberpandas/parser.py

@@ -47,8 +47,17 @@ def _to_ip_array(values):
 
     if isinstance(values, IPArray):
         return values.data
-    if not (isinstance(values, np.ndarray) and
-            values.dtype == IPType._record_type):
+
+    if (isinstance(values, np.ndarray) and
+            values.ndim == 1 and
+            np.issubdtype(values.dtype, np.integer)):
+        # We assume we're given the low bits here.
+        values = values.astype("u8")
+        values = np.asarray(values, dtype=IPType._record_type)
+        values['hi'] = 0
+
+    elif not (isinstance(values, np.ndarray) and
+              values.dtype == IPType._record_type):
         values = _to_int_pairs(values)
     return np.atleast_1d(np.asarray(values, dtype=IPType._record_type))
 
@@ -61,9 +70,6 @@ def _to_int_pairs(values):
         if values.ndim != 2:
             raise ValueError("'values' should be a 2-D when passing a "
                              "NumPy array.")
-        if values.dtype != int:
-            raise ValueError("'values' should be integer dtype when "
-                             "passing a NumPy array.")
     elif isinstance(values, tuple) and len(values) == 2:
         # like IPArray((0, 0))
         # which isn't IPArray([0, 0])

tests/ip/test_ip.py

@@ -301,3 +301,29 @@ def test_factorize():
     uniques = uniques.astype(object)
     tm.assert_numpy_array_equal(labels, expected_labels)
     tm.assert_numpy_array_equal(uniques, expected_uniques)
+
+
+@pytest.mark.parametrize('values', [
+    [0, 1, 2],
+])
+def test_from_ndarray(values):
+    result = ip.IPArray(np.asarray(values))
+    expected = ip.IPArray(values)
+    assert result.equals(expected)
+
+
+@pytest.mark.parametrize('start, stop, step, expected', [
+    (1, 3, None, [1, 2]),
+    ('0.0.0.1', '0.0.0.3', None, [1, 2]),
+    (2**64 + 1, 2**64 + 3, None, [2**64 + 1, 2**64 + 2]),
+    ('::1:0:0:0:1', '::1:0:0:0:3', None, [2**64 + 1, 2**64 + 2]),
+    (2**64 - 1, 2**64 + 2, None, [2**64 - 1, 2**64, 2**64 + 1]),
+    ('::ffff:ffff:ffff:ffff', '::1:0:0:0:2', None,
+     [2**64 - 1, 2**64, 2**64 + 1]),
+    (1, 6, 2, [1, 3, 5]),
+    ('0.0.0.1', '0.0.0.6', '0.0.0.2', [1, 3, 5]),
+])
+def test_ip_range(start, stop, step, expected):
+    result = ip.ip_range(start, stop, step)
+    expected = ip.IPArray(expected)
+    assert result.equals(expected)