Generate 'betterproto' classes for better experience in Python.

.gitignore

@@ -4,4 +4,5 @@ build/
 dist/
 python/pkg/drift_protocol.egg-info/
 python/pkg/drift_protocol
+python/pkg/drift
 !python/pkg/drift_protocol/__init__.py.in

.pre-commit-config.yaml

@@ -16,7 +16,6 @@ repos:
       name: black
       entry: black
       language: python
-      language_version: python3.8
       args:
         - .
         - --check

python/python_tools/generator.py

@@ -82,6 +82,18 @@ def generate_proto(self, source: str, create_package=True) -> str:
             str(source),
         ]
 
+        print(protoc_command)
+        if subprocess.call(protoc_command, cwd=PROTO_SPEC_FOLDER) != 0:
+            sys.exit(-1)
+
+        print(f"Generating Better Proto {output_path}...")
+        protoc_command = [
+            self.protoc,
+            "-I.",
+            f"--python_betterproto_out={self.output_dir}",
+            str(source),
+        ]
+
         print(protoc_command)
         if subprocess.call(protoc_command, cwd=PROTO_SPEC_FOLDER) != 0:
             sys.exit(-1)

python/setup.py

@@ -47,6 +47,9 @@ def update_package_version(path: Path, version: str, protoc_version: str):
     with open(path / "__init__.py", "w") as f:
         f.write(init_content)
 
+    with open(path.parent / "drift" / "__init__.py", "w") as f:
+        f.write(init_content)
+
 
 def build_version():
     """Build dynamic version and update version in package"""
@@ -144,5 +147,8 @@ def packages(self):
     package_dir={"": "pkg"},
     packages=LazyPackageFinder(finder=partial(find_packages, where="pkg")),
     python_requires=">=3.7",
-    install_requires=[f"protobuf>={PROTOBUF_VERSION}, <=3.20.3"],
+    install_requires=[
+        f"protobuf>={PROTOBUF_VERSION}, <=3.20.3",
+        "betterproto[compiler]>=2.0.0b6",
+    ],
 )

python/tests/test_package.py

@@ -4,6 +4,7 @@
 import os
 import pkgutil
 import drift_protocol
+import drift.proto
 
 
 def test_version():
@@ -20,10 +21,27 @@ def test_modules():
     assert modules
 
 
+def test_betterproto_modules():
+    """Test tha package has modules"""
+    package_path = os.path.dirname(drift.proto.__file__)
+    modules = [name for _, name, _ in pkgutil.iter_modules([package_path])]
+
+    assert modules
+
+
 def test_import():
     """Test if no errors during the import"""
     from drift_protocol.common import (
         DriftPackage,
     )
 
     assert DriftPackage()
+
+
+def test_betterproto_import():
+    """Test if no errors during the import"""
+    from drift.proto.common import (
+        DriftPackage,
+    )
+
+    assert DriftPackage() is not None