Merge pull request #56 from IBM/ServiceStreamingDescriptor

ServiceStreamingDescriptor: Support streaming rpcs in descriptor_to_file

py_to_proto/descriptor_to_file.py

@@ -8,6 +8,7 @@
 
 # Third Party
 from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pb2
 
 ## Globals #####################################################################
 
@@ -204,8 +205,26 @@ def _service_descriptor_to_file(
  lines = []
  lines.append(f"service {service_descriptor.name} {{")
  for method in service_descriptor.methods:
+ # The MethodDescriptor protobuf representation holds fields to represent
+ # server and client streaming, but these are not exposed in the python
+ # class that wraps a MethodDescriptor. They are, however, held in the
+ # underlying C implementation in upb, so the information is retained but
+ # is only accessible when re-serializing the python object to a proto
+ # representation of the descriptor.
+ md_proto = descriptor_pb2.MethodDescriptorProto()
+ method.CopyToProto(md_proto)
+ client_streaming = "stream " if md_proto.client_streaming else ""
+ server_streaming = "stream " if md_proto.server_streaming else ""
+
  lines.append(
- f"{PROTO_FILE_INDENT}rpc {method.name}({method.input_type.full_name}) returns ({method.output_type.full_name});"
+ "{}rpc {}({}{}) returns ({}{});".format(
+ PROTO_FILE_INDENT,
+ method.name,
+ client_streaming,
+ method.input_type.full_name,
+ server_streaming,
+ method.output_type.full_name,
+ )
  )
  lines.append("}")
  return _indent_lines(indent, lines)

tests/test_descriptor_to_file.py

@@ -296,10 +296,29 @@ def test_descriptor_to_file_service_descriptor(temp_dpool):
  "service": {
  "rpcs": [
  {
- "name": "FooPredict",
+ "name": "FooPredictUnaryUnary",
  "input_type": "foo.bar.Foo",
  "output_type": "foo.bar.Foo",
- }
+ },
+ {
+ "name": "FooPredictUnaryStream",
+ "input_type": "foo.bar.Foo",
+ "output_type": "foo.bar.Foo",
+ "server_streaming": True,
+ },
+ {
+ "name": "FooPredictStreamUnary",
+ "input_type": "foo.bar.Foo",
+ "output_type": "foo.bar.Foo",
+ "client_streaming": True,
+ },
+ {
+ "name": "FooPredictStreamStream",
+ "input_type": "foo.bar.Foo",
+ "output_type": "foo.bar.Foo",
+ "client_streaming": True,
+ "server_streaming": True,
+ },
  ]
  }
  },
@@ -308,6 +327,13 @@ def test_descriptor_to_file_service_descriptor(temp_dpool):
  # TODO: type annotation fixup
  res = descriptor_to_file(service_descriptor)
  assert "service FooService {" in res
+ assert "rpc FooPredictUnaryUnary(foo.bar.Foo) returns (foo.bar.Foo)" in res
+ assert "rpc FooPredictUnaryStream(foo.bar.Foo) returns (stream foo.bar.Foo)" in res
+ assert "rpc FooPredictStreamUnary(stream foo.bar.Foo) returns (foo.bar.Foo)" in res
+ assert (
+ "rpc FooPredictStreamStream(stream foo.bar.Foo) returns (stream foo.bar.Foo)"
+ in res
+ )
 
 
 def test_descriptor_to_file_compilable_proto_with_service_descriptor(temp_dpool):