Skip to content

Commit 3a13749

Browse files
committed
Allow disabling generation of grpclib stubs
1 parent 1d54ef8 commit 3a13749

File tree

3 files changed

+26
-4
lines changed

3 files changed

+26
-4
lines changed

src/betterproto/plugin/models.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,9 +181,14 @@ def comment(self) -> str:
181181

182182

183183
@dataclass
184-
class PluginRequestCompiler:
184+
class Options:
185+
grpc_kind: str = "grpclib"
186+
185187

188+
@dataclass
189+
class PluginRequestCompiler:
186190
plugin_request_obj: plugin.CodeGeneratorRequest
191+
options: Options
187192
output_packages: Dict[str, "OutputTemplate"] = field(default_factory=dict)
188193

189194
@property

src/betterproto/plugin/parser.py

Lines changed: 13 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,24 @@ 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="):
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(
86+
plugin_request_obj=request,
87+
options=options
88+
)
7889
# Gather output packages
7990
for proto_file in request.proto_file:
8091
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)