proto: enable compiling gRPC services

tensorboard/defs/protos.bzl

@@ -12,16 +12,55 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("@com_google_protobuf//:protobuf.bzl", "py_proto_library")
+load("@com_google_protobuf//:protobuf.bzl", "proto_gen")
 
-def tb_proto_library(name, srcs=None, visibility=None, testonly=None):
-  py_proto_library(
-      name = name + "_py_pb2",
-      srcs = srcs,
-      srcs_version = "PY2AND3",
-      deps = ["@com_google_protobuf//:protobuf_python"],
-      protoc = "@com_google_protobuf//:protoc",
-      visibility = visibility,
-      default_runtime = "@com_google_protobuf//:protobuf_python",
-      testonly = testonly,
-  )
+def tb_proto_library(
+        name,
+        srcs = None,
+        visibility = None,
+        testonly = None,
+        has_services = False):
+    outs_proto = _PyOuts(srcs, grpc = False)
+    outs_grpc = _PyOuts(srcs, grpc = True) if has_services else []
+    outs_all = outs_proto + outs_grpc
+
+    proto_gen(
+        name = name + "_py_pb2_genproto",
+        srcs = srcs,
+        deps = ["@com_google_protobuf//:protobuf_python_genproto"],
+        includes = [],
+        protoc = "@com_google_protobuf//:protoc",
+        gen_py = True,
+        outs = outs_all,
+        visibility = ["//visibility:public"],
+        plugin = "//external:grpc_python_plugin" if has_services else None,
+        plugin_language = "grpc",
+    )
+
+    native.py_library(
+        name = name + "_py_pb2",
+        srcs = outs_proto,
+        imports = [],
+        srcs_version = "PY2AND3",
+        deps = ["@com_google_protobuf//:protobuf_python"],
+        testonly = testonly,
+        visibility = visibility,
+    )
+    if has_services:
+        native.py_library(
+            name = name + "_py_pb2_grpc",
+            srcs = outs_grpc,
+            imports = [],
+            srcs_version = "PY2AND3",
+            deps = [
+                name + "_py_pb2",
+                "@com_google_protobuf//:protobuf_python",
+            ],
+            testonly = testonly,
+            visibility = visibility,
+        )
+
+def _PyOuts(srcs, grpc):
+    # Adapted from @com_google_protobuf//:protobuf.bzl.
+    ext = "_pb2.py" if not grpc else "_pb2_grpc.py"
+    return [s[:-len(".proto")] + ext for s in srcs]