Clean up '_assign_entity_to_mutation'.

Addresses @dhermes' comments:
#524 (comment)
#524 (comment)
#524 (comment)

Author: tseaver

gcloud/datastore/batch.py

@@ -210,8 +210,8 @@ def put(self, entity):
         if entity.key is None:
             raise ValueError("Entity must have a key")
 
-        if _assign_entity_to_mutation(self.mutation, entity):
-            self._auto_id_entities.append(entity)
+        _assign_entity_to_mutation(
+            self.mutation, entity, self._auto_id_entities)
 
     def delete(self, key):
         """Remember a key to be deleted durring ``commit``.
@@ -273,15 +273,32 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             _BATCHES.pop()
 
 
-def _assign_entity_to_mutation(mutation_pb, entity):
-    """Helper method for ``Batch.put``."""
+def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):
+    """Copy ``entity`` into appropriate slot of ``mutation_pb``.
+
+    If ``entity.key`` is incomplete, append ``entity`` to ``auto_id_entities``
+    for later fixup during ``commit``.
+
+    Helper method for ``Batch.put``.
+
+    :type mutation_pb: :class:`gcloud.datastore.datastore_v1_pb2.Mutation`
+    :param mutation_pb; the Mutation protobuf for the batch / transaction.
+
+    :type entity: :class:`gcloud.datastore.entity.Entity`
+    :param entity; the entity being updated within the batch / transaction.
+
+    :type auto_id_entities: list of :class:`gcloud.datastore.entity.Entity`
+    :param auto_id_entities: entiites with partial keys, to be fixed up
+                              during commit.
+    """
     auto_id = entity.key.is_partial
 
     key_pb = entity.key.to_protobuf()
     key_pb = helpers._prepare_key_for_request(key_pb)
 
     if auto_id:
         insert = mutation_pb.insert_auto_id.add()
+        auto_id_entities.append(entity)
     else:
         insert = mutation_pb.upsert.add()
 
@@ -301,5 +318,3 @@ def _assign_entity_to_mutation(mutation_pb, entity):
 
             for sub_value in prop.value.list_value:
                 sub_value.indexed = False
-
-    return auto_id