Redefine method definition for sync stubs

Author: AdrienVannson

src/betterproto2_compiler/plugin/typing_compiler.py

@@ -37,6 +37,14 @@ def async_iterable(self, type_: str) -> str:
     def async_iterator(self, type_: str) -> str:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def sync_iterable(self, type_: str) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def sync_iterator(self, type_: str) -> str:
+        raise NotImplementedError
+
     @abc.abstractmethod
     def imports(self) -> builtins.dict[str, set[str] | None]:
         """
@@ -89,6 +97,14 @@ def async_iterator(self, type_: str) -> str:
         self._imports["typing"].add("AsyncIterator")
         return f"AsyncIterator[{type_}]"
 
+    def sync_iterable(self, type_: str) -> str:
+        self._imports["typing"].add("Iterable")
+        return f"Iterable[{type_}]"
+
+    def sync_iterator(self, type_: str) -> str:
+        self._imports["typing"].add("Iterator")
+        return f"Iterator[{type_}]"
+
     def imports(self) -> builtins.dict[str, set[str] | None]:
         return {k: v if v else None for k, v in self._imports.items()}
 
@@ -124,6 +140,14 @@ def async_iterable(self, type_: str) -> str:
     def async_iterator(self, type_: str) -> str:
         self._imported = True
         return f"typing.AsyncIterator[{type_}]"
+    
+    def sync_iterable(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.Iterable[{type_}]"
+
+    def sync_iterator(self, type_: str) -> str:
+        self._imported = True
+        return f"typing.Iterator[{type_}]"
 
     def imports(self) -> builtins.dict[str, set[str] | None]:
         if self._imported:
@@ -159,5 +183,13 @@ def async_iterator(self, type_: str) -> str:
         self._imports["collections.abc"].add("AsyncIterator")
         return f"AsyncIterator[{type_}]"
 
+    def async_iterable(self, type_: str) -> str:
+        self._imports["collections.abc"].add("Iterable")
+        return f"Iterable[{type_}]"
+
+    def async_iterator(self, type_: str) -> str:
+        self._imports["collections.abc"].add("Iterator")
+        return f"Iterator[{type_}]"
+
     def imports(self) -> builtins.dict[str, set[str] | None]:
         return {k: v if v else None for k, v in self._imports.items()}

src/betterproto2_compiler/templates/service_stub.py.j2

@@ -10,24 +10,7 @@ class {% block class_name %}{% endblock %}({% block inherit_from %}{% endblock %
     {% block class_content %}{% endblock %}
 
     {% for method in service.methods %}
-    async def {{ method.py_name }}(self
-        {%- if not method.client_streaming -%}
-            , {{ method.py_input_message_param }}:
-                {%- if method.is_input_msg_empty -%}
-                    "{{ method.py_input_message_type }} | None" = None
-                {%- else -%}
-                    "{{ method.py_input_message_type }}"
-                {%- endif -%}
-        {%- else -%}
-            {# Client streaming: need a request iterator instead #}
-            , {{ 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)) }}"
-        {%- endif -%}
-            ,
-            *
-            , timeout: {{ output_file.typing_compiler.optional("float") }} = None
-            , deadline: "{{ output_file.typing_compiler.optional("Deadline") }}" = None
-            , metadata: "{{ output_file.typing_compiler.optional("MetadataLike") }}" = None
-            ) -> "{% if method.server_streaming %}{{ output_file.typing_compiler.async_iterator(method.py_output_message_type ) }}{% else %}{{ method.py_output_message_type }}{% endif %}":
+    {% block method_definition scoped required %}{% endblock %}
         {% if method.comment %}
         """
         {{ method.comment | indent(8) }}

src/betterproto2_compiler/templates/service_stub_async.py.j2

@@ -3,3 +3,25 @@
 {# Class definition #}
 {% block class_name %}{{ service.py_name }}Stub{% endblock %}
 {% block inherit_from %}betterproto2.ServiceStub{% endblock %}
+
+{# Methods definition #}
+{% block method_definition %}
+    async def {{ method.py_name }}(self
+        {%- if not method.client_streaming -%}
+            , {{ method.py_input_message_param }}:
+                {%- if method.is_input_msg_empty -%}
+                    "{{ method.py_input_message_type }} | None" = None
+                {%- else -%}
+                    "{{ method.py_input_message_type }}"
+                {%- endif -%}
+        {%- else -%}
+            {# Client streaming: need a request iterator instead #}
+            , {{ 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)) }}"
+        {%- endif -%}
+            ,
+            *
+            , timeout: {{ output_file.typing_compiler.optional("float") }} = None
+            , deadline: "{{ output_file.typing_compiler.optional("Deadline") }}" = None
+            , metadata: "{{ output_file.typing_compiler.optional("MetadataLike") }}" = None
+            ) -> "{% if method.server_streaming %}{{ output_file.typing_compiler.async_iterator(method.py_output_message_type ) }}{% else %}{{ method.py_output_message_type }}{% endif %}":
+{% endblock %}

src/betterproto2_compiler/templates/service_stub_sync.py.j2

@@ -8,3 +8,25 @@
     def __init__(self, channel: grpc.Channel):
         self._channel = channel
 {% endblock %}
+
+{# Methods definition #}
+{% block method_definition %}
+    def {{ method.py_name }}(self
+        {%- if not method.client_streaming -%}
+            , {{ method.py_input_message_param }}:
+                {%- if method.is_input_msg_empty -%}
+                    "{{ method.py_input_message_type }} | None" = None
+                {%- else -%}
+                    "{{ method.py_input_message_type }}"
+                {%- endif -%}
+        {%- else -%}
+            {# Client streaming: need a request iterator instead #}
+            , {{ method.py_input_message_param }}_iterator: "{{ output_file.typing_compiler.iterable(method.py_input_message_type) }}"
+        {%- endif -%}
+            ,
+            *
+            , timeout: {{ output_file.typing_compiler.optional("float") }} = None
+            , deadline: "{{ output_file.typing_compiler.optional("Deadline") }}" = None
+            , metadata: "{{ output_file.typing_compiler.optional("MetadataLike") }}" = None
+            ) -> "{% if method.server_streaming %}{{ output_file.typing_compiler.sync_iterator(method.py_output_message_type ) }}{% else %}{{ method.py_output_message_type }}{% endif %}":
+{% endblock %}