Skip to content

Commit 614a5f2

Browse files
committed
Allow disabling generation of grpclib stubs
1 parent 1d54ef8 commit 614a5f2

File tree

3 files changed

+21
-4
lines changed

3 files changed

+21
-4
lines changed

src/betterproto/plugin/models.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,11 +179,14 @@ def comment(self) -> str:
179179
proto_file=self.proto_file, path=self.path, indent=self.comment_indent
180180
)
181181

182+
@dataclass
183+
class Options:
184+
grpc_kind: str = "grpclib"
182185

183186
@dataclass
184187
class PluginRequestCompiler:
185-
186188
plugin_request_obj: plugin.CodeGeneratorRequest
189+
options: Options
187190
output_packages: Dict[str, "OutputTemplate"] = field(default_factory=dict)
188191

189192
@property

src/betterproto/plugin/parser.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
FieldCompiler,
3030
MapEntryCompiler,
3131
MessageCompiler,
32+
Options,
3233
OneOfFieldCompiler,
3334
OutputTemplate,
3435
PluginRequestCompiler,
@@ -41,7 +42,6 @@
4142
if TYPE_CHECKING:
4243
from google.protobuf.descriptor import Descriptor
4344

44-
4545
def traverse(
4646
proto_file: FieldDescriptorProto,
4747
) -> "itertools.chain[Tuple[Union[str, EnumDescriptorProto], List[int]]]":
@@ -68,13 +68,21 @@ def _traverse(
6868
_traverse([5], proto_file.enum_type), _traverse([4], proto_file.message_type)
6969
)
7070

71+
def parse_options(plugin_options: List[str]) -> Options:
72+
options = Options()
73+
for option in plugin_options:
74+
if option.startswith("grpc-kind="):
75+
options.grpc_kind = option.split("=", 1)[1]
76+
return options
7177

7278
def generate_code(
7379
request: plugin.CodeGeneratorRequest, response: plugin.CodeGeneratorResponse
7480
) -> None:
7581
plugin_options = request.parameter.split(",") if request.parameter else []
7682

77-
request_data = PluginRequestCompiler(plugin_request_obj=request)
83+
options = parse_options(plugin_options)
84+
85+
request_data = PluginRequestCompiler(plugin_request_obj=request, options=options)
7886
# Gather output packages
7987
for proto_file in request.proto_file:
8088
if (

src/betterproto/templates/template.py.j2

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@ from typing import {% for i in output_file.typing_imports|sort %}{{ i }}{% if no
1515
{% endif %}
1616

1717
import betterproto
18+
{% if output_file.parent_request.options.grpc_kind == "grpclib" %}
1819
from betterproto.grpc.grpclib_server import ServiceBase
19-
{% if output_file.services %}
20+
{% endif %}
21+
{% if output_file.services and output_file.parent_request.options.grpc_kind == "grpclib" %}
2022
import grpclib
2123
{% endif %}
2224

@@ -68,6 +70,9 @@ class {{ message.py_name }}(betterproto.Message):
6870

6971

7072
{% endfor %}
73+
74+
{% if output_file.parent_request.options.grpc_kind == "grpclib" %}
75+
7176
{% for service in output_file.services %}
7277
class {{ service.py_name }}Stub(betterproto.ServiceStub):
7378
{% if service.comment %}
@@ -237,6 +242,7 @@ class {{ service.py_name }}Base(ServiceBase):
237242
}
238243

239244
{% endfor %}
245+
{% endif %}
240246

241247
{% for i in output_file.imports|sort %}
242248
{{ i }}

0 commit comments

Comments
 (0)