Adding Key.__eq__ and __ne__ for comparison.

dhermes · dhermes · commit 90d50639609a · 2015-01-14T19:57:49.000-08:00
diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py
@@ -335,6 +335,17 @@ def parent(self):
 
         return self._parent
 
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+
+        return (self.namespace == other.namespace and
+                self.path == other.path and
+                _dataset_ids_equal(self.dataset_id, other.dataset_id))
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
     def __repr__(self):
         return '<Key%s, dataset=%s>' % (self.path, self.dataset_id)
 
@@ -347,6 +358,14 @@ def _validate_dataset_id(dataset_id, parent):
 
     If ``dataset_id`` is unset, attempt to infer the ID from the environment.
 
+    :type dataset_id: string
+    :param dataset_id: A dataset ID.
+
+    :type parent: :class:`gcloud.datastore.key.Key` or ``NoneType``
+    :param parent: The parent of the key or ``None``.
+
+    :rtype: string
+    :returns: The ``dataset_id`` passed in, or implied from the environment.
     :raises: :class:`ValueError` if ``dataset_id`` is ``None`` and no dataset
              can be inferred.
     """
@@ -360,3 +379,55 @@ def _validate_dataset_id(dataset_id, parent):
             dataset_id = _implicit_environ.DATASET_ID
 
     return dataset_id
+
+
+def _dataset_ids_equal(dataset_id1, dataset_id2):
+    """Compares two dataset IDs for fuzzy equality.
+
+    Each may be prefixed or unprefixed (but not null, since dataset ID
+    is required on a key). The only allowed prefixes are 's~' and 'e~'.
+
+    Two identical prefixed match
+
+      >>> 's~foo' == 's~foo'
+      >>> 'e~bar' == 'e~bar'
+
+    while non-identical prefixed don't
+
+      >>> 's~foo' != 's~bar'
+      >>> 's~foo' != 'e~foo'
+
+    As for non-prefixed, they can match other non-prefixed or
+    prefixed:
+
+      >>> 'foo' == 'foo'
+      >>> 'foo' == 's~foo'
+      >>> 'foo' == 'e~foo'
+      >>> 'foo' != 'bar'
+      >>> 'foo' != 's~bar'
+
+    (Ties are resolved since 'foo' can only be an alias for one of
+    s~foo or e~foo in the backend.)
+
+    :type dataset_id1: string
+    :param dataset_id1: A dataset ID.
+
+    :type dataset_id2: string
+    :param dataset_id2: A dataset ID.
+
+    :rtype: boolean
+    :returns: Boolean indicating if the IDs are the same.
+    """
+    if dataset_id1 == dataset_id2:
+        return True
+
+    if dataset_id1.startswith('s~') or dataset_id1.startswith('e~'):
+        # If `dataset_id1` is prefixed and not matching, then the only way
+        # they can match is if `dataset_id2` is unprefixed.
+        return dataset_id1[2:] == dataset_id2
+    elif dataset_id2.startswith('s~') or dataset_id2.startswith('e~'):
+        # Here we know `dataset_id1` is unprefixed and `dataset_id2`
+        # is prefixed.
+        return dataset_id1 == dataset_id2[2:]
+
+    return False
diff --git a/gcloud/datastore/test_key.py b/gcloud/datastore/test_key.py
@@ -20,7 +20,6 @@ class TestKey(unittest2.TestCase):
     _DEFAULT_DATASET = 'DATASET'
 
     def setUp(self):
-
         from gcloud.datastore import _implicit_environ
         self._replaced_dataset_id = _implicit_environ.DATASET_ID
         _implicit_environ.DATASET_ID = None
@@ -291,3 +290,58 @@ def test_parent_multiple_calls(self):
         self.assertEqual(parent.path, _PARENT_PATH)
         new_parent = key.parent
         self.assertTrue(parent is new_parent)
+
+    def test___eq__bad_type(self):
+        key = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        self.assertNotEqual(key, None)
+        self.assertNotEqual(key, object())
+        self.assertNotEqual(None, key)
+
+    def test___eq__same_value(self):
+        key1 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        key2 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        self.assertFalse(key1 is key2)
+        self.assertEqual(key1, key2)
+
+    def test___eq__different_dataset(self):
+        ALT_DATASET = 'DATASET-ALT'
+        key1 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        key2 = self._makeOne('KIND', dataset_id=ALT_DATASET)
+
+        self.assertNotEqual(ALT_DATASET, self._DEFAULT_DATASET)
+        self.assertNotEqual(key1, key2)
+
+    def test___eq__different_path(self):
+        key1 = self._makeOne('KIND', 10, dataset_id=self._DEFAULT_DATASET)
+        key2 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        self.assertNotEqual(key1, key2)
+
+    def test___eq__different_namespace(self):
+        key1 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET,
+                             namespace='NAMESPACE')
+        key2 = self._makeOne('KIND', dataset_id=self._DEFAULT_DATASET)
+        self.assertNotEqual(key1, key2)
+
+
+class Test__dataset_ids_equal(unittest2.TestCase):
+
+    def _callFUT(self, dataset_id1, dataset_id2):
+        from gcloud.datastore.key import _dataset_ids_equal
+        return _dataset_ids_equal(dataset_id1, dataset_id2)
+
+    def test_identical_prefixed(self):
+        self.assertTrue(self._callFUT('s~foo', 's~foo'))
+        self.assertTrue(self._callFUT('e~bar', 'e~bar'))
+
+    def test_different_prefixed(self):
+        self.assertFalse(self._callFUT('s~foo', 's~bar'))
+        self.assertFalse(self._callFUT('s~foo', 'e~foo'))
+
+    def test_all_unprefixed(self):
+        self.assertTrue(self._callFUT('foo', 'foo'))
+        self.assertFalse(self._callFUT('foo', 'bar'))
+
+    def test_unprefixed_with_prefixed(self):
+        self.assertTrue(self._callFUT('foo', 's~foo'))
+        self.assertTrue(self._callFUT('foo', 'e~foo'))
+        self.assertFalse(self._callFUT('foo', 's~bar'))
diff --git a/regression/datastore.py b/regression/datastore.py
@@ -88,10 +88,7 @@ def _generic_test_post(self, name=None, key_id=None):
             self.assertEqual(entity.key.id, key_id)
         retrieved_entity, = datastore.get([entity.key])
         # Check the keys are the same.
-        self.assertEqual(retrieved_entity.key.path, entity.key.path)
-        self.assertEqual(retrieved_entity.key.namespace, entity.key.namespace)
-        self.assertTrue(_compare_dataset_ids(
-            retrieved_entity.key.dataset_id, entity.key.dataset_id))
+        self.assertEqual(retrieved_entity.key, entity.key)
 
         # Check the data is the same.
         retrieved_dict = dict(retrieved_entity.items())
@@ -157,8 +154,7 @@ def test_save_key_self_reference(self):
 
         stored_person = stored_persons[0]
         self.assertEqual(stored_person['fullName'], entity['fullName'])
-        self.assertEqual(stored_person.key.path, key.path)
-        self.assertEqual(stored_person.key.namespace, key.namespace)
+        self.assertEqual(stored_person.key, key)
 
 
 class TestDatastoreQuery(TestDatastore):
@@ -263,14 +259,7 @@ def test_projection_query(self):
                          {'name': 'Catelyn', 'family': 'Tully'})
 
         # Check both Catelyn keys are the same.
-        catelyn_stark_key = catelyn_stark_entity.key
-        catelyn_tully_key = catelyn_tully_entity.key
-        self.assertEqual(catelyn_stark_key.path, catelyn_tully_key.path)
-        self.assertEqual(catelyn_stark_key.namespace,
-                         catelyn_tully_key.namespace)
-        # Also check the dataset_id since both retrieved from datastore.
-        self.assertEqual(catelyn_stark_key.dataset_id,
-                         catelyn_tully_key.dataset_id)
+        self.assertEqual(catelyn_stark_entity.key, catelyn_tully_entity.key)
 
         sansa_entity = entities[8]
         sansa_dict = dict(sansa_entity.items())
@@ -352,19 +341,3 @@ def test_transaction(self):
         retrieved_dict = dict(retrieved_entity.items())
         entity_dict = dict(entity.items())
         self.assertEqual(retrieved_dict, entity_dict)
-
-
-def _compare_dataset_ids(dataset_id1, dataset_id2):
-    if dataset_id1 == dataset_id2:
-        return True
-
-    if dataset_id1.startswith('s~') or dataset_id1.startswith('e~'):
-        # If `dataset_id1` is prefixed and not matching, then the only way
-        # they can match is if `dataset_id2` is unprefixed.
-        return dataset_id1[2:] == dataset_id2
-    elif dataset_id2.startswith('s~') or dataset_id2.startswith('e~'):
-        # Here we know `dataset_id1` is unprefixed and `dataset_id2`
-        # is prefixed.
-        return dataset_id1 == dataset_id2[2:]
-
-    return False
