collapse_address_list is too slow

ipaddr.py

 #!/usr/bin/python
 #
 # Copyright 2007 Google Inc.
 #  Licensed to PSF under a Contributor Agreement.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #      http://www.apache.org/licenses/LICENSE-2.0
@@ skipping 151 matching lines @@
 
 def IPv6(address, host=False):
     """Generic IPv6 Factory function..
 
     Return an IPv6 address or network.
     """
     if host:
         return IPv6Address(address)
     return IPv6Network(address)
 
+
+def _find_address_range(addresses):
+    """Find a sequence of addresses.
+
+    Args:
+        addresses: a list of IPv4 or IPv6 addresses.
+
+    Returns:
+        A tuple containing the first and last IP addresses in the sequence.
+    
+    """
+    first = last = addresses[0]
+    for ip in addresses[1:]:
+        if ip._ip == last._ip + 1:
+            last = ip
+        else:
+            break
+    return (first, last)
+
+def _get_prefix_length(number1, number2, bits):
+    """Get the number of leading bits that are same for two numbers.
+    
+    Args:
+        number1: an integer.
+        number2: another integer.
+        bits: the maximum number of bits to compare.
+
+    Returns:
+        The number of leading bits that are the same for two numbers.
+            
+    """
+    for i in range(bits):
+        if number1 >> i == number2 >> i:
+            return bits - i
+    return 0
+
+def _count_righthand_zero_bits(number, bits):
+    """Count the number of zero bits on the right hand side.
+    
+    Args:
+        number: an integer.
+        bits: maximum number of bits to count.
+    
+    Returns:
+        The number of zero bits on the right hand side of the number.
+    
+    """
+    if number == 0:
+        return bits
+    for i in range(bits):
+        if (number >> i) % 2:
+            return i
+
+def summarize_address_range(first, last):
+    """Summarize a network range given the first and last IP addresses.
+
+    Example:
+        >>> summarize_address_range(IPv4Address('1.1.1.0'),
+            IPv4Address('1.1.1.130'))
+        [IPv4Network('1.1.1.0/25'), IPv4Network('1.1.1.128/31'),
+        IPv4Network('1.1.1.130/32')]
+
+    Args:
+        first: the first IPv4 or IPv6 address in the range.
+        last: the last IPv4 or IPv6 address in the range.
+
+    Returns:
+        The address range collapsed to a list of IPv4 or IPv6 networks.
+
+    Raise:
+        IPTypeError:
+            If the first and last objects are not IP addresses.
+            If the first and last objects are not the same version.
+        ValueError:
+            If the last object is not greater than the first.
+            If the version is not 4 or 6.
+
+    """
+
+    if not (isinstance(first, BaseIP) and isinstance(last, BaseIP)):
+        raise IPTypeError("first and last must be IP addresses, not networks")
+    if first.version != last.version:
+        raise IPTypeError("IP addresses must be same version")
+    if first > last:
+        raise ValueError("last IP address must be greater than first")
+
+    networks = []
+
+    if first.version == 4:
+        ip = IPv4Network
+    elif first.version == 6:
+        ip = IPv6Network
+    else:
+        raise ValueError("unknown IP version")
+
+    ip_bits = first._max_prefixlen
+    first = first._ip
+    last = last._ip
+    while first <= last:
+        nbits = _count_righthand_zero_bits(first, ip_bits)
+        current = None
+        while nbits >= 0:
+            addend = 2**nbits - 1
+            current = first + addend
+            nbits -= 1
+            if current <= last:
+                break
+        prefix = _get_prefix_length(first, current, ip_bits)
+        net = ip(first)
+        net.prefixlen = prefix
+        networks.append(net)
+        if current == ip._ALL_ONES:
+            break
+        first = current + 1
+    return networks
 
 def _collapse_address_list_recursive(addresses):
     """Loops through the addresses, collapsing concurrent netblocks.
 
     Example:
 
         ip1 = IPv4('1.1.0.0/24')
         ip2 = IPv4('1.1.1.0/24')
         ip3 = IPv4('1.1.2.0/24')
         ip4 = IPv4('1.1.3.0/24')
@@ skipping 41 matching lines @@
         collapse_address_list([IPv4('1.1.0.0/24'), IPv4('1.1.1.0/24')]) ->
           [IPv4('1.1.0.0/23')]
 
     Args:
         addresses: A list of IPv4 or IPv6 objects.
 
     Returns:
         A list of IPv4 or IPv6 objects depending on what we were passed.
 
     """
+    i = 0
+    addrs = []
+    ips = []
+    nets = []
+    # split IP addresses and networks
+    for ip in addresses:
+        if isinstance(ip, BaseIP):
+            ips.append(ip)
+        elif ip._prefixlen == ip._max_prefixlen:
+            ips.append(ip.ip)
+        else:
+            nets.append(ip)
+    # sort and dedup
+    ips = sorted(set(ips))
+    nets = sorted(set(nets))
+        
+    while i < len(ips):
+        (first, last) = _find_address_range(ips[i:])
+        i = ips.index(last) + 1
+        addrs.extend(summarize_address_range(first, last))
+
     return _collapse_address_list_recursive(
-        sorted(addresses, key=BaseNet._get_networks_key))
+        sorted(addrs + nets, key=BaseNet._get_networks_key))
 
 # backwards compatibility
 CollapseAddrList = collapse_address_list
 
 # Test whether this Python implementation supports byte objects that
 # are not identical to str ones.
 # We need to exclude platforms where bytes == str so that we can
 # distinguish between packed representations and strings, for example
 # b'12::' (the IPv4 address 49.50.58.58) and '12::' (an IPv6 address).
 try:
@@ skipping 74 matching lines @@
             return self._version > other._version
         if self._ip != other._ip:
             return self._ip > other._ip
         return False
 
     def __repr__(self):
         return '%s(%r)' % (self.__class__.__name__, str(self))
 
     def __str__(self):
         return  '%s' % self._string_from_ip_int(self._ip)
+
+    def __hash__(self):
+        return hash(self._ip)
 
     @property
     def version(self):
         raise NotImplementedError('BaseIP has no version')
 
 
 class BaseNet(IPAddrBase):
 
     """A generic IP object.
 
@@ skipping 348 matching lines @@
     The following methods are used by IPv4 objects in both single IP
     addresses and networks.
 
     """
 
     # Equivalent to 255.255.255.255 or 32 bits of 1's.
     _ALL_ONES = (2**32) - 1
 
     def __init__(self, address):
         self._version = 4
+        self._max_prefixlen = 32
 
     def _explode_shorthand_ip_string(self, ip_str=None):
         if not ip_str:
             ip_str = str(self)
         return ip_str
 
     def _ip_int_from_string(self, ip_str):
         """Turn the given IP string into an integer for comparison.
 
         Args:
@@ skipping 46 matching lines @@
         if len(octets) != 4:
             return False
 
         for octet in octets:
             try:
                 if not 0 <= int(octet) <= 255:
                     return False
             except ValueError:
                 return False
         return True
+
+    @property
+    def max_prefixlen(self):
+        return self._max_prefixlen
 
     @property
     def packed(self):
         """The binary representation of this address."""
         return struct.pack('!I', self._ip)
 
     @property
     def version(self):
         return self._version
 
@@ skipping 334 matching lines @@
 
     The following methods are used by IPv6 objects in both single IP
     addresses and networks.
 
     """
 
     _ALL_ONES = (2**128) - 1
 
     def __init__(self, address):
         self._version = 6
+        self._max_prefixlen = 128
 
     def _ip_int_from_string(self, ip_str=None):
         """Turn an IPv6 ip_str into an integer.
 
         Args:
             ip_str: A string, the IPv6 ip_str.
 
         Returns:
             A long, the IPv6 ip_str.
 
@@ skipping 186 matching lines @@
         Args:
             ip_str: A string, the IPv6 address.
 
         Returns:
             A boolean, True if the address is shortened.
 
         """
         if ip_str.count('::') == 1:
             return True
         return False
+
+    @property
+    def max_prefixlen(self):
+        return self._max_prefixlen
 
     @property
     def packed(self):
         """The binary representation of this address."""
         return struct.pack('!QQ', self._ip >> 64, self._ip & (2**64 - 1))
 
     @property
     def version(self):
         return self._version
 
@@ skipping 310 matching lines @@
         """
         try:
             prefixlen = int(prefixlen)
         except ValueError:
             return False
         return 0 <= prefixlen <= 128
 
     # backwards compatibility
     Subnet = subnet
     Supernet = supernet