Auto-generate GRPC client bindings as part of pip install (#272)

.github/workflows/python-client-release-to-pypi.yml

@@ -19,8 +19,8 @@ jobs:
           repo-token: ${{ secrets.GITHUB_TOKEN }}
       - uses: ./.github/workflows/python-tests
         with:
-          python-version: '3.8'
-          tox-env: 'py38'
+          python-version: '3.9'
+          tox-env: 'py39'
           path: 'client/python'
           github-token: ${{secrets.GITHUB_TOKEN}}
       - name: Publish package to PyPI

.github/workflows/python-client.yml

@@ -34,13 +34,16 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        python: [ '3.8', '3.9', '3.10' ]
+        python: [ '3.9', '3.10', '3.11', '3.12' ]
         include:
-          - tox-env: 'py38'
           - tox-env: 'py39'
             python: '3.9'
           - tox-env: 'py310'
             python: '3.10'
+          - tox-env: 'py311'
+            python: '3.11'
+          - tox-env: 'py312'
+            python: '3.12'
     steps:
       - uses: actions/checkout@v4
       - name: Setup Go

.github/workflows/python-tests/action.yml

@@ -45,7 +45,7 @@ runs:
       working-directory: ${{ inputs.path }}
     - name: Build and verify wheel
       run: |
-        python -m build --wheel
+        python -m build --sdist
         twine check dist/*
       shell: bash
       working-directory: ${{ inputs.path }}

.gitignore

@@ -77,6 +77,7 @@ client/python/dist
 *_pb2.py
 *_pb2.pyi
 *_pb2_grpc.py
+client/python/armada_client/proto/
 client/python/armada_client/armada/
 .tox
 proto-airflow

client/python/armada_client/gen/event_typings.py

@@ -1,3 +1,5 @@
+import argparse
+from pathlib import Path
 import sys
 
 from armada_client.armada import event_pb2, submit_pb2
@@ -63,15 +65,7 @@ def gen_file(states, classes, jobstates):
     return import_text, states_text, union_text, jobstates_text
 
 
-def write_file(import_text, states_text, union_text, jobstates_text, file):
-    with open(f"{file}", "w", encoding="utf-8") as f:
-        f.write(import_text)
-        f.write(states_text)
-        f.write(jobstates_text)
-        f.write(union_text)
-
-
-def main():
+def main(typings_file: Path):
     states = get_event_states()
     print("Done creating EventStates")
 
@@ -84,13 +78,16 @@ def main():
     import_text, states_text, union_text, jobstates_text = gen_file(
         states, classes, jobstates
     )
-    write_file(import_text, states_text, union_text, jobstates_text, typings_file)
+    typings_file.write_text(import_text + states_text + jobstates_text + union_text)
 
 
 if __name__ == "__main__":
-    # get path to this files location
-    root = f"{sys.path[0]}/../../"
-    typings_file = f"{root}/armada_client/typings.py"
-
-    main()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("typings_file", type=Path, help="Path to typings file")
+
+    args = parser.parse_args()
+    print(f"{args}")
+    typings_file = args.typings_file or Path("armada_client") / "typings.py"
+    print(f"{typings_file}")
+    main(typings_file)
     sys.exit(0)

client/python/pyproject.toml

@@ -1,21 +1,21 @@
 [project]
 name = "armada_client"
-version = "0.3.5"
+version = "0.4.5"
 description = "Armada gRPC API python client"
 readme = "README.md"
-requires-python = ">=3.7"
-dependencies = ["grpcio==1.66.1", "grpcio-tools==1.66.1", "mypy-protobuf>=3.2.0", "protobuf>=5.26.1,<6.0dev" ]
+requires-python = ">=3.9"
+dependencies = ["grpcio-tools", "protobuf>3.20,<5.0"]
 license = { text = "Apache Software License" }
 authors = [{ name = "G-Research Open Source Software", email = "[email protected]" }]
 
 [project.optional-dependencies]
 format = ["black==23.7.0", "flake8==7.0.0", "pylint==2.17.5"]
 # note(JayF): sphinx-jekyll-builder was broken by sphinx-markdown-builder 0.6 -- so pin to 0.5.5
 docs = ["sphinx==7.1.2", "sphinx-jekyll-builder==0.3.0", "sphinx-toolbox==3.2.0b1", "sphinx-markdown-builder==0.5.5"]
-test = ["pytest==7.3.1", "coverage>=6.5.0", "pytest-asyncio==0.21.1"]
+test = ["pytest==7.3.1", "pytest-cov", "pytest-asyncio==0.21.1"]
 
 [build-system]
-requires = ["setuptools"]
+requires = ["setuptools", "wheel", "grpcio-tools", "mypy-protobuf", "protobuf>3.20,<5.0"]
 build-backend = "setuptools.build_meta"
 
 [tool.mypy]
@@ -39,4 +39,4 @@ omit = [
 # py.typed is required for mypy to find type hints in the package
 # from: https://mypy.readthedocs.io/en/stable/installed_packages.html#making-pep-561-compatible-packages
 [tool.setuptools.package-data]
-"*" = ["*.pyi", "py.typed"]
+"*" = ["*.pyi", "py.typed", "proto/**/*.proto"]

client/python/setup.py

@@ -0,0 +1,124 @@
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import sys
+from typing import Dict
+from setuptools import setup
+import importlib.resources
+import re
+from setuptools.command.build_py import build_py
+
+
+def generate_grpc_bindings(build_lib: Path):
+    import grpc_tools.protoc
+
+    proto_include = importlib.resources.path("grpc_tools", "_proto")
+    proto_files = [
+        "google/api/annotations.proto",
+        "google/api/http.proto",
+        "github.com/gogo/protobuf/gogoproto/gogo.proto",
+        "k8s.io/api/core/v1/generated.proto",
+        "k8s.io/apimachinery/pkg/api/resource/generated.proto",
+        "k8s.io/apimachinery/pkg/apis/meta/v1/generated.proto",
+        "k8s.io/apimachinery/pkg/runtime/generated.proto",
+        "k8s.io/apimachinery/pkg/runtime/schema/generated.proto",
+        "k8s.io/apimachinery/pkg/util/intstr/generated.proto",
+        "k8s.io/api/networking/v1/generated.proto",
+        "armada/event.proto",
+        "armada/submit.proto",
+        "armada/health.proto",
+        "armada/job.proto",
+        "armada/binoculars.proto",
+    ]
+    target_root = build_lib.absolute() / "armada_client"
+
+    for proto_file in proto_files:
+        command = [
+            f"-I{proto_include}",
+            f"-I{target_root / 'proto'}",
+            f"--python_out={target_root}",
+            f"--grpc_python_out={target_root}",
+            f"--mypy_out={target_root}",
+            str(target_root / "proto" / proto_file),
+        ]
+        if grpc_tools.protoc.main(command) != 0:
+            raise Exception(f"grpc_tools.protoc.main: {command} failed")
+
+    shutil.rmtree(target_root / "github.com")
+    shutil.rmtree(target_root / "k8s.io")
+
+    adjust_import_paths(target_root)
+
+
+def adjust_import_paths(output_dir: Path):
+    replacements = {
+        r"from armada": "from armada_client.armada",
+        r"from github.com": "from armada_client.github.com",
+        r"from google.api": "from armada_client.google.api",
+    }
+
+    for file in output_dir.glob("armada/*.py"):
+        replace_in_file(file, replacements)
+    for file in output_dir.glob("google/api/*.py"):
+        replace_in_file(file, replacements)
+
+    replacements = {
+        r"from k8s.io": "from armada_client.k8s.io",
+    }
+    for file in output_dir.glob("../**/*.py"):
+        replace_in_file(file, replacements)
+
+    replacements = {
+        r" k8s": " armada_client.k8s",
+        r"\[k8s": "[armada_client.k8s",
+        r"import k8s.io": "import armada_client.k8s.io",
+    }
+    for file in output_dir.glob("k8s/**/*.pyi"):
+        replace_in_file(file, replacements)
+
+
+def replace_in_file(file: Path, replacements: Dict[str, str]):
+    """Replace patterns in a file based on the replacements dictionary."""
+
+    content = file.read_text()
+    for pattern, replacement in replacements.items():
+        content = re.sub(pattern, replacement, content)
+    file.write_text(content)
+
+
+def generate_typings(build_dir: Path):
+    typings = build_dir.absolute() / "armada_client" / "typings.py"
+    result = subprocess.run(
+        args=[
+            sys.executable,
+            str(build_dir.absolute() / "armada_client" / "gen" / "event_typings.py"),
+            str(typings),
+        ],
+        env={"PYTHONPATH": str(build_dir.absolute())},
+        capture_output=True,
+    )
+    if result.returncode != 0:
+        print(result.stdout)
+        print(result.stderr)
+    result.check_returncode()
+
+
+class BuildPackageProtos(build_py):
+    """
+    Generate GRPC code before building the package.
+    """
+
+    def run(self):
+        super().run()
+        output_dir = Path(".") if self.editable_mode else Path(self.build_lib)
+        generate_grpc_bindings(output_dir)
+        generate_typings(output_dir)
+
+
+setup(
+    cmdclass={
+        "build_py": BuildPackageProtos,
+        "develop": BuildPackageProtos,
+    },
+)

client/python/tox.ini

@@ -1,16 +1,15 @@
 [tox]
-isolated_build = true
 envlist =
     format
-    py38
     py39
     py310
+    py311
+    py312
 
 [testenv]
 extras = test
 commands =
-    coverage run -m pytest tests/unit/
-    coverage xml
+    pytest --cov={envsitepackagesdir}/armada_client --cov-report=xml --cov-report=term tests/unit/
 
 [testenv:docs]
 extras = docs

scripts/build-python-client.sh

@@ -5,37 +5,4 @@
 mkdir -p proto/armada
 cp pkg/api/event.proto pkg/api/submit.proto pkg/api/health.proto pkg/api/job.proto pkg/api/binoculars/binoculars.proto proto/armada
 sed -i 's/\([^\/]\)pkg\/api/\1armada/g' proto/armada/*.proto
-
-# generate python stubs
-cd proto
-python3 -m grpc_tools.protoc -I. --plugin=protoc-gen-mypy=$(which protoc-gen-mypy) --python_out=../client/python/armada_client --grpc_python_out=../client/python/armada_client --mypy_out=../client/python/armada_client \
-    google/api/annotations.proto \
-    google/api/http.proto \
-    armada/event.proto armada/submit.proto armada/health.proto armada/job.proto armada/binoculars.proto \
-    github.com/gogo/protobuf/gogoproto/gogo.proto \
-    k8s.io/api/core/v1/generated.proto \
-    k8s.io/apimachinery/pkg/api/resource/generated.proto \
-    k8s.io/apimachinery/pkg/apis/meta/v1/generated.proto \
-    k8s.io/apimachinery/pkg/runtime/generated.proto \
-    k8s.io/apimachinery/pkg/runtime/schema/generated.proto \
-    k8s.io/apimachinery/pkg/util/intstr/generated.proto \
-    k8s.io/api/networking/v1/generated.proto
-
-cd ..
-# This hideous code is because we can't use python package option in grpc.
-# See https://github.com/protocolbuffers/protobuf/issues/7061 for an explanation.
-# We need to import these packages as a module.
-sed -i 's/from armada/from armada_client.armada/g' client/python/armada_client/armada/*.py
-sed -i 's/from github.com/from armada_client.github.com/g' client/python/armada_client/armada/*.py
-sed -i 's/from google.api/from armada_client.google.api/g' client/python/armada_client/armada/*.py
-sed -i 's/from google.api/from armada_client.google.api/g' client/python/armada_client/google/api/*.py
-
-find client/python/armada_client/ -name '*.py' | xargs sed -i 's/from k8s.io/from armada_client.k8s.io/g'
-
-# Generate better docs for the client
-export PYTHONPATH=${PWD}/client/python
-python3 ${PWD}/client/python/armada_client/gen/event_typings.py
-
-find client/python/armada_client/k8s -name '*.pyi' | xargs sed -i 's/ k8s/ armada_client.k8s/g'
-find client/python/armada_client/k8s -name '*.pyi' | xargs sed -i 's/\[k8s/\[armada_client.k8s/g'
-find client/python/armada_client/k8s/io -name '*.pyi' | xargs sed -i 's/import k8s.io/import armada_client.k8s.io/g'
+cp -rf proto/* client/python/armada_client/proto/