+ make networks and addresses not sortable by default.

 + provide a key function suitable for passing to sorted
   enabling the sorting.

git-svn-id: https://ipaddr-py.googlecode.com/svn@140 09200d28-7f98-11dd-ad27-0f66e57d2035

2 files changed, 68 insertions, 9 deletions

diff --git a/trunk/ipaddr.py b/trunk/ipaddr.py
index 3e14445..7e7cb0c 100644
--- a/trunk/ipaddr.py
+++ b/trunk/ipaddr.py
@@ -345,6 +345,29 @@ try:
 except NameError: # <Python2.6
     _compat_has_real_bytes = False
 
+def get_mixed_type_key(obj):
+    """Return a key suitable for sorting between networks and addresses.
+
+    Address and Network objects are not sortable by default; they're
+    fundamentally different so the expression
+
+        IPv4Address('1.1.1.1') <= IPv4Network('1.1.1.1/24')
+
+    doesn't make any sense.  There are some times however, where you may wish
+    to have ipaddr sort these for you anyway. If you need to do this, you
+    can use this function as the key= argument to sorted().
+    
+    Args:
+      obj: either a Network or Address object.
+    Returns:
+      appropriate key.
+
+    """
+    if isinstance(obj, _BaseNet):
+        return obj._get_networks_key()
+    elif isinstance(obj, _BaseIP):
+        return obj._get_address_key()
+    return NotImplemented
 
 class _IPAddrBase(object):
 
@@ -410,16 +433,22 @@ class _BaseIP(_IPAddrBase):
 
     def __lt__(self, other):
         if self._version != other._version:
-            raise TypeError("%s and %s are not of the same version" % (
-                str(self), str(other)))        
+            raise TypeError('%s and %s are not of the same version' % (
+                    str(self), str(other)))
+        if not isinstance(other, _BaseIP):
+            raise TypeError('%s and %s are not of the same type' % (
+                    str(self), str(other)))
         if self._ip != other._ip:
             return self._ip < other._ip
         return False
 
     def __gt__(self, other):
         if self._version != other._version:
-            raise TypeError("%s and %s are not of the same version" % (
-                str(self), str(other)))        
+            raise TypeError('%s and %s are not of the same version' % (
+                    str(self), str(other)))
+        if not isinstance(other, _BaseIP):
+            raise TypeError('%s and %s are not of the same type' % (
+                    str(self), str(other)))
         if self._ip != other._ip:
             return self._ip > other._ip
         return False
@@ -433,6 +462,9 @@ class _BaseIP(_IPAddrBase):
     def __hash__(self):
         return hash(hex(self._ip))
 
+    def _get_address_key(self):
+        return (self._version, self)
+
     @property
     def version(self):
         raise NotImplementedError('BaseIP has no version')
@@ -488,8 +520,11 @@ class _BaseNet(_IPAddrBase):
 
     def __lt__(self, other):
         if self._version != other._version:
-            raise TypeError("%s and %s are not of the same version" % (
-                str(self), str(other)))        
+            raise TypeError('%s and %s are not of the same version' % (
+                    str(self), str(other)))
+        if not isinstance(other, _BaseNet):
+            raise TypeError('%s and %s are not of the same type' % (
+                    str(self), str(other)))
         if self.network != other.network:
             return self.network < other.network
         if self.netmask != other.netmask:
@@ -498,8 +533,11 @@ class _BaseNet(_IPAddrBase):
 
     def __gt__(self, other):
         if self._version != other._version:
-            raise TypeError("%s and %s are not of the same version" % (
-                str(self), str(other)))        
+            raise TypeError('%s and %s are not of the same version' % (
+                    str(self), str(other)))
+        if not isinstance(other, _BaseNet):
+            raise TypeError('%s and %s are not of the same type' % (
+                    str(self), str(other)))
         if self.network != other.network:
             return self.network > other.network
         if self.netmask != other.netmask:
diff --git a/trunk/ipaddr_test.py b/trunk/ipaddr_test.py
index f4ce79d..33fc886 100755
--- a/trunk/ipaddr_test.py
+++ b/trunk/ipaddr_test.py
@@ -100,7 +100,6 @@ class IpaddrUnitTest(unittest.TestCase):
         self.assertEqual(str(self.ipv6.hostmask),
                          '::ffff:ffff:ffff:ffff')
 
-
     def testBadVersionComparison(self):
         # These should always raise TypeError
         v4addr = ipaddr.IPAddress('1.1.1.1')
@@ -118,6 +117,28 @@ class IpaddrUnitTest(unittest.TestCase):
         self.assertRaises(TypeError, v6net.__lt__, v4net)
         self.assertRaises(TypeError, v6net.__gt__, v4net)
 
+    def testMixedTypeComparison(self):
+        v4addr = ipaddr.IPAddress('1.1.1.1')
+        v4net = ipaddr.IPNetwork('1.1.1.1/32')
+        v6addr = ipaddr.IPAddress('::1')
+        v6net = ipaddr.IPNetwork('::1/128')
+        
+        self.assertRaises(TypeError, v4addr.__lt__, v4net)
+        self.assertRaises(TypeError, v4addr.__gt__, v4net)
+        self.assertRaises(TypeError, v4net.__lt__, v4addr)
+        self.assertRaises(TypeError, v4net.__gt__, v4addr)
+
+        self.assertRaises(TypeError, v6addr.__lt__, v6net)
+        self.assertRaises(TypeError, v6addr.__gt__, v6net)
+        self.assertRaises(TypeError, v6net.__lt__, v6addr)
+        self.assertRaises(TypeError, v6net.__gt__, v6addr)
+
+        # with get_mixed_type_key, you can sort addresses and network.
+        self.assertEqual([v4addr, v4net], sorted([v4net, v4addr],
+                                                 key=ipaddr.get_mixed_type_key))
+        self.assertEqual([v6addr, v6net], sorted([v6net, v6addr],
+                                                 key=ipaddr.get_mixed_type_key))
+        
     def testIpFromInt(self):
         self.assertEqual(self.ipv4.ip, ipaddr.IPv4Network(16909060).ip)
         self.assertRaises(ipaddr.AddressValueError,