Fix #533: add key and entity equality methods

gcloud/datastore/batch.py

@@ -21,6 +21,7 @@ class Local(object):
 
 from gcloud.datastore import _implicit_environ
 from gcloud.datastore import helpers
+from gcloud.datastore.key import _dataset_ids_equal
 from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
 
 
@@ -216,8 +217,7 @@ def put(self, entity):
         if entity.key is None:
             raise ValueError("Entity must have a key")
 
-        if not helpers._dataset_ids_equal(self._dataset_id,
-                                          entity.key.dataset_id):
+        if not _dataset_ids_equal(self._dataset_id, entity.key.dataset_id):
             raise ValueError("Key must be from same dataset as batch")
 
         _assign_entity_to_mutation(
@@ -235,8 +235,7 @@ def delete(self, key):
         if key.is_partial:
             raise ValueError("Key must be complete")
 
-        if not helpers._dataset_ids_equal(self._dataset_id,
-                                          key.dataset_id):
+        if not _dataset_ids_equal(self._dataset_id, key.dataset_id):
             raise ValueError("Key must be from same dataset as batch")
 
         key_pb = key.to_protobuf()

gcloud/datastore/entity.py

@@ -77,6 +77,32 @@ def __init__(self, key=None, exclude_from_indexes=()):
         self.key = key
         self._exclude_from_indexes = set(exclude_from_indexes)
 
+    def __eq__(self, other):
+        """Compare two entities for equality.
+
+        Entities compare equal if their keys compare equal, and their
+        properties compare equal.
+
+        :rtype: boolean
+        :returns: True if the entities compare equal, else False.
+        """
+        if not isinstance(other, Entity):
+            raise NotImplementedError
+
+        return (self.key == other.key and
+                dict(self) == dict(other))
+
+    def __ne__(self, other):
+        """Compare two entities for inequality.
+
+        Entities compare equal if their keys compare equal, and their
+        properties compare equal.
+
+        :rtype: boolean
+        :returns: False if the entities compare equal, else True.
+        """
+        return not self == other
+
     @property
     def kind(self):
         """Get the kind of the current entity.

gcloud/datastore/helpers.py

@@ -319,55 +319,3 @@ def _add_keys_to_request(request_field_pb, key_pbs):
     for key_pb in key_pbs:
         key_pb = _prepare_key_for_request(key_pb)
         request_field_pb.add().CopyFrom(key_pb)
-
-
-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

gcloud/datastore/key.py

@@ -75,6 +75,50 @@ def __init__(self, *path_args, **kwargs):
         # _combine_args() is called.
         self._path = self._combine_args()
 
+    def __eq__(self, other):
+        """Compare two keys for equality.
+
+        Incomplete keys never compare equal to any other key.
+
+        Completed keys compare equal if they have the same path, dataset ID,
+        and namespace.
+
+        :rtype: boolean
+        :returns: True if the keys compare equal, else False.
+        """
+        if not isinstance(other, Key):
+            raise NotImplementedError
+
+        if self.is_partial or other.is_partial:
+            return False
+
+        return (self.flat_path == other.flat_path and
+                _dataset_ids_equal(self.dataset_id, other.dataset_id) and
+                self.namespace == other.namespace)
+
+    def __ne__(self, other):
+        """Compare two keys for inequality.
+
+        Incomplete keys never compare equal to any other key.
+
+        Completed keys compare equal if they have the same path, dataset ID,
+        and namespace.
+
+        :rtype: boolean
+        :returns: False if the keys compare equal, else True.
+        """
+        return not self == other
+
+    def __hash__(self):
+        """Hash a keys for use in a dictionary lookp.
+
+        :rtype: integer
+        :returns: a hash of the key's state.
+        """
+        return (hash(self.flat_path) +
+                hash(self.dataset_id) +
+                hash(self.namespace))
+
     @staticmethod
     def _parse_path(path_args):
         """Parses positional arguments into key path with kinds and IDs.
@@ -362,3 +406,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

gcloud/datastore/test_entity.py

@@ -53,6 +53,44 @@ def test_ctor_explicit(self):
         self.assertEqual(sorted(entity.exclude_from_indexes),
                          sorted(_EXCLUDE_FROM_INDEXES))
 
+    def test___eq_____ne___w_non_entitye(self):
+        from gcloud.datastore.key import Key
+        key = Key(_KIND, _ID, dataset_id=_DATASET_ID)
+        entity = self._makeOne(key=key)
+        self.assertRaises(NotImplementedError, lambda: entity == object())
+        self.assertRaises(NotImplementedError, lambda: entity != object())
+
+    def test___eq_____ne___w_different_keys(self):
+        from gcloud.datastore.key import Key
+        _ID1 = 1234
+        _ID2 = 2345
+        key1 = Key(_KIND, _ID1, dataset_id=_DATASET_ID)
+        entity1 = self._makeOne(key=key1)
+        key2 = Key(_KIND, _ID2, dataset_id=_DATASET_ID)
+        entity2 = self._makeOne(key=key2)
+        self.assertFalse(entity1 == entity2)
+        self.assertTrue(entity1 != entity2)
+
+    def test___eq_____ne___w_same_keys(self):
+        from gcloud.datastore.key import Key
+        key1 = Key(_KIND, _ID, dataset_id=_DATASET_ID)
+        entity1 = self._makeOne(key=key1)
+        key2 = Key(_KIND, _ID, dataset_id=_DATASET_ID)
+        entity2 = self._makeOne(key=key2)
+        self.assertTrue(entity1 == entity2)
+        self.assertFalse(entity1 != entity2)
+
+    def test___eq_____ne___w_same_keys_different_props(self):
+        from gcloud.datastore.key import Key
+        key1 = Key(_KIND, _ID, dataset_id=_DATASET_ID)
+        entity1 = self._makeOne(key=key1)
+        entity1['foo'] = 'Foo'
+        key2 = Key(_KIND, _ID, dataset_id=_DATASET_ID)
+        entity2 = self._makeOne(key=key2)
+        entity1['bar'] = 'Bar'
+        self.assertFalse(entity1 == entity2)
+        self.assertTrue(entity1 != entity2)
+
     def test___repr___no_key_empty(self):
         entity = self._makeOne()
         self.assertEqual(repr(entity), '<Entity {}>')

gcloud/datastore/test_helpers.py

@@ -525,27 +525,3 @@ def test_prepare_dataset_id_unset(self):
         key = datastore_pb.Key()
         new_key = self._callFUT(key)
         self.assertTrue(new_key is key)
-
-
-class Test__dataset_ids_equal(unittest2.TestCase):
-
-    def _callFUT(self, dataset_id1, dataset_id2):
-        from gcloud.datastore.helpers 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'))