Remove '_get_dataset_id_from_keys'.

Author: tseaver

gcloud/datastore/api.py

@@ -60,31 +60,6 @@ def _require_connection(connection=None):
     return connection
 
 
-def _get_dataset_id_from_keys(keys):
-    """Determines dataset ID from a list of keys.
-
-    :type keys: list of :class:`gcloud.datastore.key.Key`
-    :param keys: The keys from the same dataset.
-
-    :rtype: string
-    :returns: The dataset ID of the keys.
-    :raises: :class:`ValueError` if the key dataset IDs don't agree.
-    """
-    if any(key is None for key in keys):
-        raise ValueError('None not allowed')
-
-    dataset_id = keys[0].dataset_id
-    # Rather than creating a list or set of all dataset IDs, we iterate
-    # and check. We could allow the backend to check this for us if IDs
-    # with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59)
-    # or if we made sure that a prefix s~ or e~ was on each key.
-    for key in keys[1:]:
-        if key.dataset_id != dataset_id:
-            raise ValueError('All keys in get must be from the same dataset.')
-
-    return dataset_id
-
-
 def get(keys, missing=None, deferred=None, connection=None):
     """Retrieves entities, along with their attributes.
 
@@ -111,7 +86,7 @@ def get(keys, missing=None, deferred=None, connection=None):
         return []
 
     connection = _require_connection(connection)
-    dataset_id = _get_dataset_id_from_keys(keys)
+    dataset_id, = set([key.dataset_id for key in keys])
 
     transaction = Transaction.current()
 
@@ -157,8 +132,7 @@ def put(entities, connection=None):
     in_batch = current is not None
     if not in_batch:
         keys = [entity.key for entity in entities]
-        dataset_id = _get_dataset_id_from_keys(keys)
-        current = Batch(dataset_id=dataset_id, connection=connection)
+        current = Batch(dataset_id=keys[0].dataset_id, connection=connection)
     for entity in entities:
         current.put(entity)
     if not in_batch:
@@ -183,8 +157,7 @@ def delete(keys, connection=None):
     current = Batch.current()
     in_batch = current is not None
     if not in_batch:
-        dataset_id = _get_dataset_id_from_keys(keys)
-        current = Batch(dataset_id=dataset_id, connection=connection)
+        current = Batch(dataset_id=keys[0].dataset_id, connection=connection)
     for key in keys:
         current.delete(key)
     if not in_batch:

gcloud/datastore/test_api.py

@@ -90,37 +90,6 @@ def test_implicit_set_passed_explicitly(self):
             self.assertTrue(self._callFUT(CONNECTION) is CONNECTION)
 
 
-class Test__get_dataset_id_from_keys(unittest2.TestCase):
-
-    def _callFUT(self, keys):
-        from gcloud.datastore.api import _get_dataset_id_from_keys
-        return _get_dataset_id_from_keys(keys)
-
-    def _make_key(self, dataset_id):
-
-        class _Key(object):
-            def __init__(self, dataset_id):
-                self.dataset_id = dataset_id
-
-        return _Key(dataset_id)
-
-    def test_empty(self):
-        self.assertRaises(IndexError, self._callFUT, [])
-
-    def test_w_None(self):
-        self.assertRaises(ValueError, self._callFUT, [None])
-
-    def test_w_mismatch(self):
-        key1 = self._make_key('foo')
-        key2 = self._make_key('bar')
-        self.assertRaises(ValueError, self._callFUT, [key1, key2])
-
-    def test_w_match(self):
-        key1 = self._make_key('foo')
-        key2 = self._make_key('foo')
-        self.assertEqual(self._callFUT([key1, key2]), 'foo')
-
-
 class Test_get_function(unittest2.TestCase):
 
     def _callFUT(self, keys, missing=None, deferred=None, connection=None):