Updating datastore script to add imported message types.

dhermes · dhermes · commit cd9bc263895a · 2016-08-17T15:01:00.000-07:00
diff --git a/gcloud/datastore/_generated/datastore_grpc_pb2.py b/gcloud/datastore/_generated/datastore_grpc_pb2.py
@@ -1,10 +1,4 @@
-import grpc
-from grpc.beta import implementations as beta_implementations
-from grpc.beta import interfaces as beta_interfaces
-from grpc.framework.common import cardinality
-from grpc.framework.interfaces.face import utilities as face_utilities
-
-# BEGIN: Manually added imports
+# BEGIN: Imports from datastore_pb2
 from gcloud.datastore._generated.datastore_pb2 import AllocateIdsRequest
 from gcloud.datastore._generated.datastore_pb2 import AllocateIdsResponse
 from gcloud.datastore._generated.datastore_pb2 import BeginTransactionRequest
@@ -13,11 +7,19 @@
 from gcloud.datastore._generated.datastore_pb2 import CommitResponse
 from gcloud.datastore._generated.datastore_pb2 import LookupRequest
 from gcloud.datastore._generated.datastore_pb2 import LookupResponse
+from gcloud.datastore._generated.datastore_pb2 import Mutation
+from gcloud.datastore._generated.datastore_pb2 import MutationResult
+from gcloud.datastore._generated.datastore_pb2 import ReadOptions
 from gcloud.datastore._generated.datastore_pb2 import RollbackRequest
 from gcloud.datastore._generated.datastore_pb2 import RollbackResponse
 from gcloud.datastore._generated.datastore_pb2 import RunQueryRequest
 from gcloud.datastore._generated.datastore_pb2 import RunQueryResponse
-#   END: Manually added imports
+#   END: Imports from datastore_pb2
+import grpc
+from grpc.beta import implementations as beta_implementations
+from grpc.beta import interfaces as beta_interfaces
+from grpc.framework.common import cardinality
+from grpc.framework.interfaces.face import utilities as face_utilities
 
 
 class DatastoreStub(object):
diff --git a/scripts/make_datastore_grpc.py b/scripts/make_datastore_grpc.py
@@ -29,6 +29,8 @@
 GRPC_ONLY_FILE = os.path.join(ROOT_DIR, 'gcloud', 'datastore',
                               '_generated', 'datastore_grpc_pb2.py')
 GRPCIO_VIRTUALENV = os.environ.get('GRPCIO_VIRTUALENV', 'protoc')
+MESSAGE_SNIPPET = ' = _reflection.GeneratedProtocolMessageType('
+IMPORT_TEMPLATE = 'from gcloud.datastore._generated.datastore_pb2 import %s\n'
 
 
 def get_pb2_contents_with_grpc():
@@ -110,10 +112,32 @@ def get_pb2_grpc_only():
     return grpc_only_lines
 
 
+def get_pb2_message_types():
+    """Get message types defined in datastore pb2 file.
+
+    :rtype: list
+    :returns: A list of names that are defined as message types.
+    """
+    non_grpc_contents = get_pb2_contents_without_grpc()
+    result = []
+    for line in non_grpc_contents:
+        if MESSAGE_SNIPPET in line:
+            name, _ = line.split(MESSAGE_SNIPPET)
+            result.append(name)
+
+    return sorted(result)
+
+
 def main():
     """Write gRPC-only lines to custom module."""
     grpc_only_lines = get_pb2_grpc_only()
     with open(GRPC_ONLY_FILE, 'wb') as file_obj:
+        # First add imports for public objects in the original.
+        file_obj.write('# BEGIN: Imports from datastore_pb2\n')
+        for name in get_pb2_message_types():
+            import_line = IMPORT_TEMPLATE % (name,)
+            file_obj.write(import_line)
+        file_obj.write('#   END: Imports from datastore_pb2\n')
         file_obj.write(''.join(grpc_only_lines))
 
 
