Add support for pydantic dataclasses (#406)

tests/generate.py

@@ -11,6 +11,7 @@ from tests.util import (
     get_directories,
     inputs_path,
     output_path_betterproto,
+    output_path_betterproto_pydantic,
     output_path_reference,
     protoc,
 )
@@ -80,9 +81,11 @@ async def generate_test_case_output(
 
     test_case_output_path_reference = output_path_reference.joinpath(test_case_name)
     test_case_output_path_betterproto = output_path_betterproto
+    test_case_output_path_betterproto_pyd = output_path_betterproto_pydantic
 
     os.makedirs(test_case_output_path_reference, exist_ok=True)
     os.makedirs(test_case_output_path_betterproto, exist_ok=True)
+    os.makedirs(test_case_output_path_betterproto_pyd, exist_ok=True)
 
     clear_directory(test_case_output_path_reference)
     clear_directory(test_case_output_path_betterproto)
@@ -90,9 +93,13 @@ async def generate_test_case_output(
     (
         (ref_out, ref_err, ref_code),
         (plg_out, plg_err, plg_code),
+        (plg_out_pyd, plg_err_pyd, plg_code_pyd),
     ) = await asyncio.gather(
         protoc(test_case_input_path, test_case_output_path_reference, True),
         protoc(test_case_input_path, test_case_output_path_betterproto, False),
+        protoc(
+            test_case_input_path, test_case_output_path_betterproto_pyd, False, True
+        ),
     )
 
     if ref_code == 0:
@@ -131,7 +138,27 @@ async def generate_test_case_output(
             sys.stderr.buffer.write(plg_err)
             sys.stderr.buffer.flush()
 
-    return max(ref_code, plg_code)
+    if plg_code_pyd == 0:
+        print(
+            f"\033[31;1;4mGenerated plugin (pydantic compatible) output for {test_case_name!r}\033[0m"
+        )
+    else:
+        print(
+            f"\033[31;1;4mFailed to generate plugin (pydantic compatible) output for {test_case_name!r}\033[0m"
+        )
+
+    if verbose:
+        if plg_out_pyd:
+            print("Plugin stdout:")
+            sys.stdout.buffer.write(plg_out_pyd)
+            sys.stdout.buffer.flush()
+
+        if plg_err_pyd:
+            print("Plugin stderr:")
+            sys.stderr.buffer.write(plg_err_pyd)
+            sys.stderr.buffer.flush()
+
+    return max(ref_code, plg_code, plg_code_pyd)
 
 
 HELP = "\n".join(

tests/inputs/bool/test_bool.py

@@ -1,6 +1,19 @@
+import pytest
+
 from tests.output_betterproto.bool import Test
+from tests.output_betterproto_pydantic.bool import Test as TestPyd
 
 
 def test_value():
     message = Test()
     assert not message.value, "Boolean is False by default"
+
+
+def test_pydantic_no_value():
+    with pytest.raises(ValueError):
+        TestPyd()
+
+
+def test_pydantic_value():
+    message = Test(value=False)
+    assert not message.value

tests/inputs/oneof/test_oneof.py

@@ -1,5 +1,6 @@
 import betterproto
 from tests.output_betterproto.oneof import Test
+from tests.output_betterproto_pydantic.oneof import Test as TestPyd
 from tests.util import get_test_case_json_data
 
 
@@ -13,3 +14,8 @@ def test_which_name():
     message = Test()
     message.from_json(get_test_case_json_data("oneof", "oneof_name.json")[0].json)
     assert betterproto.which_one_of(message, "foo") == ("pitier", "Mr. T")
+
+
+def test_which_count_pyd():
+    message = TestPyd(pitier="Mr. T", just_a_regular_field=2, bar_name="a_bar")
+    assert betterproto.which_one_of(message, "foo") == ("pitier", "Mr. T")

tests/util.py

@@ -1,7 +1,10 @@
 import asyncio
+import atexit
 import importlib
 import os
+import platform
 import sys
+import tempfile
 from dataclasses import dataclass
 from pathlib import Path
 from types import ModuleType
@@ -22,6 +25,7 @@ root_path = Path(__file__).resolve().parent
 inputs_path = root_path.joinpath("inputs")
 output_path_reference = root_path.joinpath("output_reference")
 output_path_betterproto = root_path.joinpath("output_betterproto")
+output_path_betterproto_pydantic = root_path.joinpath("output_betterproto_pydantic")
 
 
 def get_files(path, suffix: str) -> Generator[str, None, None]:
@@ -36,19 +40,56 @@ def get_directories(path):
 
 
 async def protoc(
-    path: Union[str, Path], output_dir: Union[str, Path], reference: bool = False
+    path: Union[str, Path],
+    output_dir: Union[str, Path],
+    reference: bool = False,
+    pydantic_dataclasses: bool = False,
 ):
     path: Path = Path(path).resolve()
     output_dir: Path = Path(output_dir).resolve()
     python_out_option: str = "python_betterproto_out" if not reference else "python_out"
-    command = [
-        sys.executable,
-        "-m",
-        "grpc.tools.protoc",
-        f"--proto_path={path.as_posix()}",
-        f"--{python_out_option}={output_dir.as_posix()}",
-        *[p.as_posix() for p in path.glob("*.proto")],
-    ]
+
+    if pydantic_dataclasses:
+        plugin_path = Path("src/betterproto/plugin/main.py")
+
+        if "Win" in platform.system():
+            with tempfile.NamedTemporaryFile(
+                "w", encoding="UTF-8", suffix=".bat", delete=False
+            ) as tf:
+                # See https://stackoverflow.com/a/42622705
+                tf.writelines(
+                    [
+                        "@echo off",
+                        f"\nchdir {os.getcwd()}",
+                        f"\n{sys.executable} -u {plugin_path.as_posix()}",
+                    ]
+                )
+
+                tf.flush()
+
+                plugin_path = Path(tf.name)
+                atexit.register(os.remove, plugin_path)
+
+        command = [
+            sys.executable,
+            "-m",
+            "grpc.tools.protoc",
+            f"--plugin=protoc-gen-custom={plugin_path.as_posix()}",
+            "--experimental_allow_proto3_optional",
+            "--custom_opt=pydantic_dataclasses",
+            f"--proto_path={path.as_posix()}",
+            f"--custom_out={output_dir.as_posix()}",
+            *[p.as_posix() for p in path.glob("*.proto")],
+        ]
+    else:
+        command = [
+            sys.executable,
+            "-m",
+            "grpc.tools.protoc",
+            f"--proto_path={path.as_posix()}",
+            f"--{python_out_option}={output_dir.as_posix()}",
+            *[p.as_posix() for p in path.glob("*.proto")],
+        ]
     proc = await asyncio.create_subprocess_exec(
         *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
     )