Skip to content
This repository was archived by the owner on Jun 9, 2025. It is now read-only.

Commit 52ab890

Browse files
committed
Use jinja templates to define service stubs, add sync stub
1 parent 9c5782e commit 52ab890

File tree

4 files changed

+105
-96
lines changed

4 files changed

+105
-96
lines changed
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
class {% block class_name %}{% endblock %}(betterproto2.ServiceStub):
2+
{% if service.comment %}
3+
"""
4+
{{ service.comment | indent(4) }}
5+
"""
6+
{% elif not service.methods %}
7+
pass
8+
{% endif %}
9+
10+
{% for method in service.methods %}
11+
async def {{ method.py_name }}(self
12+
{%- if not method.client_streaming -%}
13+
, message:
14+
{%- if method.is_input_msg_empty -%}
15+
"{{ method.py_input_message_type }} | None" = None
16+
{%- else -%}
17+
"{{ method.py_input_message_type }}"
18+
{%- endif -%}
19+
{%- else -%}
20+
{# Client streaming: need a request iterator instead #}
21+
, messages: "AsyncIterable[{{ method.py_input_message_type }}] | Iterable[{{ method.py_input_message_type }}]"
22+
{%- endif -%}
23+
,
24+
*
25+
, timeout: "float | None" = None
26+
, deadline: "Deadline | None" = None
27+
, metadata: "MetadataLike | None" = None
28+
) -> "{% if method.server_streaming %}AsyncIterator[{{ method.py_output_message_type }}]{% else %}{{ method.py_output_message_type }}{% endif %}":
29+
{% if method.comment %}
30+
"""
31+
{{ method.comment | indent(8) }}
32+
"""
33+
{% endif %}
34+
35+
{% if method.deprecated %}
36+
warnings.warn("{{ service.py_name }}.{{ method.py_name }} is deprecated", DeprecationWarning)
37+
{% endif %}
38+
39+
{% if method.server_streaming %}
40+
{% if method.client_streaming %}
41+
async for response in self._stream_stream(
42+
"{{ method.route }}",
43+
messages,
44+
{{ method.py_input_message_type }},
45+
{{ method.py_output_message_type }},
46+
timeout=timeout,
47+
deadline=deadline,
48+
metadata=metadata,
49+
):
50+
yield response
51+
{% else %}{# i.e. not client streaming #}
52+
{% if method.is_input_msg_empty %}
53+
if message is None:
54+
message = {{ method.py_input_message_type }}()
55+
56+
{% endif %}
57+
async for response in self._unary_stream(
58+
"{{ method.route }}",
59+
message,
60+
{{ method.py_output_message_type }},
61+
timeout=timeout,
62+
deadline=deadline,
63+
metadata=metadata,
64+
):
65+
yield response
66+
67+
{% endif %}{# if client streaming #}
68+
{% else %}{# i.e. not server streaming #}
69+
{% if method.client_streaming %}
70+
return await self._stream_unary(
71+
"{{ method.route }}",
72+
messages,
73+
{{ method.py_input_message_type }},
74+
{{ method.py_output_message_type }},
75+
timeout=timeout,
76+
deadline=deadline,
77+
metadata=metadata,
78+
)
79+
{% else %}{# i.e. not client streaming #}
80+
{% if method.is_input_msg_empty %}
81+
if message is None:
82+
message = {{ method.py_input_message_type }}()
83+
84+
{% endif %}
85+
return await self._unary_unary(
86+
"{{ method.route }}",
87+
message,
88+
{{ method.py_output_message_type }},
89+
timeout=timeout,
90+
deadline=deadline,
91+
metadata=metadata,
92+
)
93+
{% endif %}{# client streaming #}
94+
{% endif %}
95+
96+
{% endfor %}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
{% extends "service_stub.py.j2" %}
2+
3+
{% block class_name %}{{ service.py_name }}Stub{% endblock %}
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
{% extends "service_stub.py.j2" %}
2+
3+
{% block class_name %}{{ service.py_name }}SyncStub{% endblock %}

src/betterproto2_compiler/templates/template.py.j2

Lines changed: 3 additions & 96 deletions
Original file line numberDiff line numberDiff line change
@@ -85,103 +85,10 @@ default_message_pool.register_message("{{ output_file.package }}", "{{ message.p
8585

8686

8787
{% endfor %}
88-
{% for _, service in output_file.services|dictsort(by="key") %}
89-
class {{ service.py_name }}Stub(betterproto2.ServiceStub):
90-
{% if service.comment %}
91-
"""
92-
{{ service.comment | indent(4) }}
93-
"""
94-
{% elif not service.methods %}
95-
pass
96-
{% endif %}
97-
98-
{% for method in service.methods %}
99-
async def {{ method.py_name }}(self
100-
{%- if not method.client_streaming -%}
101-
, message:
102-
{%- if method.is_input_msg_empty -%}
103-
"{{ method.py_input_message_type }} | None" = None
104-
{%- else -%}
105-
"{{ method.py_input_message_type }}"
106-
{%- endif -%}
107-
{%- else -%}
108-
{# Client streaming: need a request iterator instead #}
109-
, messages: "AsyncIterable[{{ method.py_input_message_type }}] | Iterable[{{ method.py_input_message_type }}]"
110-
{%- endif -%}
111-
,
112-
*
113-
, timeout: "float | None" = None
114-
, deadline: "Deadline | None" = None
115-
, metadata: "MetadataLike | None" = None
116-
) -> "{% if method.server_streaming %}AsyncIterator[{{ method.py_output_message_type }}]{% else %}{{ method.py_output_message_type }}{% endif %}":
117-
{% if method.comment %}
118-
"""
119-
{{ method.comment | indent(8) }}
120-
"""
121-
{% endif %}
12288

123-
{% if method.deprecated %}
124-
warnings.warn("{{ service.py_name }}.{{ method.py_name }} is deprecated", DeprecationWarning)
125-
{% endif %}
126-
127-
{% if method.server_streaming %}
128-
{% if method.client_streaming %}
129-
async for response in self._stream_stream(
130-
"{{ method.route }}",
131-
messages,
132-
{{ method.py_input_message_type }},
133-
{{ method.py_output_message_type }},
134-
timeout=timeout,
135-
deadline=deadline,
136-
metadata=metadata,
137-
):
138-
yield response
139-
{% else %}{# i.e. not client streaming #}
140-
{% if method.is_input_msg_empty %}
141-
if message is None:
142-
message = {{ method.py_input_message_type }}()
143-
144-
{% endif %}
145-
async for response in self._unary_stream(
146-
"{{ method.route }}",
147-
message,
148-
{{ method.py_output_message_type }},
149-
timeout=timeout,
150-
deadline=deadline,
151-
metadata=metadata,
152-
):
153-
yield response
154-
155-
{% endif %}{# if client streaming #}
156-
{% else %}{# i.e. not server streaming #}
157-
{% if method.client_streaming %}
158-
return await self._stream_unary(
159-
"{{ method.route }}",
160-
messages,
161-
{{ method.py_input_message_type }},
162-
{{ method.py_output_message_type }},
163-
timeout=timeout,
164-
deadline=deadline,
165-
metadata=metadata,
166-
)
167-
{% else %}{# i.e. not client streaming #}
168-
{% if method.is_input_msg_empty %}
169-
if message is None:
170-
message = {{ method.py_input_message_type }}()
171-
172-
{% endif %}
173-
return await self._unary_unary(
174-
"{{ method.route }}",
175-
message,
176-
{{ method.py_output_message_type }},
177-
timeout=timeout,
178-
deadline=deadline,
179-
metadata=metadata,
180-
)
181-
{% endif %}{# client streaming #}
182-
{% endif %}
183-
184-
{% endfor %}
89+
{% for _, service in output_file.services|dictsort(by="key") %}
90+
{% include "service_stub_async.py.j2" %}
91+
{% include "service_stub_sync.py.j2" %}
18592
{% endfor %}
18693

18794
{% for i in output_file.imports_end %}

0 commit comments

Comments
 (0)