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

Commit d2e1e0a

Browse files
AdrienVannsonAdrienVannson
committed
Use jinja templates to define service stubs, add sync stub
1 parent 5863013 commit d2e1e0a

File tree

4 files changed

+105
-95
lines changed

4 files changed

+105
-95
lines changed

‎src/betterproto2_compiler/templates/service_stub.py.j2

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+
, {{ method.py_input_message_param }}:
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+
, {{ method.py_input_message_param }}_iterator: "{{ output_file.typing_compiler.union(output_file.typing_compiler.async_iterable(method.py_input_message_type), output_file.typing_compiler.iterable(method.py_input_message_type)) }}"
22+
{%- endif -%}
23+
,
24+
*
25+
, timeout: {{ output_file.typing_compiler.optional("float") }} = None
26+
, deadline: "{{ output_file.typing_compiler.optional("Deadline") }}" = None
27+
, metadata: "{{ output_file.typing_compiler.optional("MetadataLike") }}" = None
28+
) -> "{% if method.server_streaming %}{{ output_file.typing_compiler.async_iterator(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.proto_obj.options and method.proto_obj.options.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+
{{ method.py_input_message_param }}_iterator,
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 {{ method.py_input_message_param }} is None:
54+
{{ method.py_input_message_param }} = {{ method.py_input_message_type }}()
55+
56+
{% endif %}
57+
async for response in self._unary_stream(
58+
"{{ method.route }}",
59+
{{ method.py_input_message_param }},
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+
{{ method.py_input_message_param }}_iterator,
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 {{ method.py_input_message_param }} is None:
82+
{{ method.py_input_message_param }} = {{ method.py_input_message_type }}()
83+
84+
{% endif %}
85+
return await self._unary_unary(
86+
"{{ method.route }}",
87+
{{ method.py_input_message_param }},
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 %}

‎src/betterproto2_compiler/templates/service_stub_async.py.j2

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 %}

‎src/betterproto2_compiler/templates/service_stub_sync.py.j2

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 & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -70,103 +70,11 @@ class {{ message.py_name }}(betterproto2.Message):
7070
{% endif %}
7171

7272
{% endfor %}
73-
{% for _, service in output_file.services|dictsort(by="key") %}
74-
class {{ service.py_name }}Stub(betterproto2.ServiceStub):
75-
{% if service.comment %}
76-
"""
77-
{{ service.comment | indent(4) }}
78-
"""
79-
{% elif not service.methods %}
80-
pass
81-
{% endif %}
82-
83-
{% for method in service.methods %}
84-
async def {{ method.py_name }}(self
85-
{%- if not method.client_streaming -%}
86-
, {{ method.py_input_message_param }}:
87-
{%- if method.is_input_msg_empty -%}
88-
"{{ method.py_input_message_type }} | None" = None
89-
{%- else -%}
90-
"{{ method.py_input_message_type }}"
91-
{%- endif -%}
92-
{%- else -%}
93-
{# Client streaming: need a request iterator instead #}
94-
, {{ method.py_input_message_param }}_iterator: "{{ output_file.typing_compiler.union(output_file.typing_compiler.async_iterable(method.py_input_message_type), output_file.typing_compiler.iterable(method.py_input_message_type)) }}"
95-
{%- endif -%}
96-
,
97-
*
98-
, timeout: {{ output_file.typing_compiler.optional("float") }} = None
99-
, deadline: "{{ output_file.typing_compiler.optional("Deadline") }}" = None
100-
, metadata: "{{ output_file.typing_compiler.optional("MetadataLike") }}" = None
101-
) -> "{% if method.server_streaming %}{{ output_file.typing_compiler.async_iterator(method.py_output_message_type ) }}{% else %}{{ method.py_output_message_type }}{% endif %}":
102-
{% if method.comment %}
103-
"""
104-
{{ method.comment | indent(8) }}
105-
"""
106-
{% endif %}
107-
108-
{% if method.proto_obj.options and method.proto_obj.options.deprecated %}
109-
warnings.warn("{{ service.py_name }}.{{ method.py_name }} is deprecated", DeprecationWarning)
110-
{% endif %}
111-
112-
{% if method.server_streaming %}
113-
{% if method.client_streaming %}
114-
async for response in self._stream_stream(
115-
"{{ method.route }}",
116-
{{ method.py_input_message_param }}_iterator,
117-
{{ method.py_input_message_type }},
118-
{{ method.py_output_message_type }},
119-
timeout=timeout,
120-
deadline=deadline,
121-
metadata=metadata,
122-
):
123-
yield response
124-
{% else %}{# i.e. not client streaming #}
125-
{% if method.is_input_msg_empty %}
126-
if {{ method.py_input_message_param }} is None:
127-
{{ method.py_input_message_param }} = {{ method.py_input_message_type }}()
128-
129-
{% endif %}
130-
async for response in self._unary_stream(
131-
"{{ method.route }}",
132-
{{ method.py_input_message_param }},
133-
{{ method.py_output_message_type }},
134-
timeout=timeout,
135-
deadline=deadline,
136-
metadata=metadata,
137-
):
138-
yield response
139-
140-
{% endif %}{# if client streaming #}
141-
{% else %}{# i.e. not server streaming #}
142-
{% if method.client_streaming %}
143-
return await self._stream_unary(
144-
"{{ method.route }}",
145-
{{ method.py_input_message_param }}_iterator,
146-
{{ method.py_input_message_type }},
147-
{{ method.py_output_message_type }},
148-
timeout=timeout,
149-
deadline=deadline,
150-
metadata=metadata,
151-
)
152-
{% else %}{# i.e. not client streaming #}
153-
{% if method.is_input_msg_empty %}
154-
if {{ method.py_input_message_param }} is None:
155-
{{ method.py_input_message_param }} = {{ method.py_input_message_type }}()
15673

157-
{% endif %}
158-
return await self._unary_unary(
159-
"{{ method.route }}",
160-
{{ method.py_input_message_param }},
161-
{{ method.py_output_message_type }},
162-
timeout=timeout,
163-
deadline=deadline,
164-
metadata=metadata,
165-
)
166-
{% endif %}{# client streaming #}
167-
{% endif %}
74+
{% for _, service in output_file.services|dictsort(by="key") %}
75+
{% include "service_stub_async.py.j2" %}
16876

169-
{% endfor %}
77+
{% include "service_stub_sync.py.j2" %}
17078
{% endfor %}
17179

17280
{% for i in output_file.imports_end %}

0 commit comments

Comments
 (0)