Add .pyi debug symbols to python wheel (opensearch-project#253)

* Migrate top level BUILD.bazel to python_proto_library.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Move py output back to py_library rule and fix deps/protos usage in python_proto_library.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* protos mandatory for python_proto_library.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Migrate service/schema python to compile rules for additional flags and options.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Fix imports.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Remove python generation in schema/service - Aggregate all at top level.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Remove py libraries from schema/service dirs - Now aggregated in top level.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Incremental bazelmod for rules-proto-grpc.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Bump bazel -> 7.3.0 for compatibility as well as C++ -> 17.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Refactor python wheel genrule due to libraries aggregated at top level.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Try explicit file names.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Bazel may not allow wildcards in genrules.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Fix path in pywheel genrule.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Include seperate target for .pyi files in genrule deps.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Pass genrule artifacts directly to py_wheel rule.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Rules library to python 3.10.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Follow symlink to fix sed failure.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Add gRPC dependencies back in to wheel with .pyi symbols.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Formatting + comments.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Move to rules_proto_grpc 5.2.0. Still support protobuf 3.x.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Revert bazelmod migration - Move everything back to WORKSPACE and fix genrule paths.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Add protoc genrule to create pyi files and copy into wheel.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Changelog.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Remove unused diff.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Update docker + docs with new py commands.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

* Combine python schema/service commands in docs/ci/docker.

Signed-off-by: Finn Carroll <carrofin@amazon.com>

---------

Signed-off-by: Finn Carroll <carrofin@amazon.com>

Author: finnegancarroll

.github/workflows/build-protobufs-python.yml

@@ -15,4 +15,4 @@ jobs:
           distribution: temurin # Temurin is a distribution of adoptium
           java-version: 21
       - name: build protobufs python
-        run: bazel build //:python_protos_all
+        run: bazel build //:python_schemas //:python_services

BUILD.bazel

@@ -1,11 +1,10 @@
 load("@io_bazel_rules_go//go:def.bzl", "go_library")
 load("@versioning//:version.bzl", "VERSION")
+load("@rules_proto_grpc//python:defs.bzl", "python_proto_library", "python_grpc_library")
+load("@rules_python//python:packaging.bzl", "py_wheel")
 
 package(default_visibility = ["//visibility:public"])
 
-load("@rules_python//python:defs.bzl", "py_library")
-load("@rules_python//python:packaging.bzl", "py_wheel")
-
 java_library(
     name = "java_protos_all",
     visibility = ["//visibility:public"],
@@ -18,18 +17,28 @@ java_library(
     ],
 )
 
-py_library(
-    name = "python_protos_all",
+python_proto_library(
+    name = "python_schemas",
     visibility = ["//visibility:public"],
-    deps = [
-        "//protos/schemas:common_python_proto",
-        "//protos/schemas:document_python_proto",
-        "//protos/schemas:search_python_proto",
-        "//protos/services:document_service_python_proto",
-        "//protos/services:search_service_python_proto",
+    protos = [
+        "//protos/schemas:common_proto",
+        "//protos/schemas:document_proto",
+        "//protos/schemas:search_proto",
     ],
 )
 
+python_grpc_library(
+    name = "python_services",
+    visibility = ["//visibility:public"],
+    protos = [
+        "//protos/services:document_service_proto",
+        "//protos/services:search_service_proto",
+    ],
+    deps = [
+        ":python_schemas"
+    ]
+)
+
 go_library(
     name = "go_protos_all",
     visibility = ["//visibility:public"],
@@ -44,96 +53,85 @@ go_library(
 
 """
 Package python libraries into a wheel.
-Here we seperate schema and service libraries between two py_library definitions and execute
-genrules on both to fix the imports. Protoc does not provide a package option for python so we must fix
-the import statements in these proto files manually. See issues:
+Protoc does not provide a package option for python so we must fix the import statements in these proto files manually:
 https://github.com/protocolbuffers/protobuf/issues/7061
 https://github.com/protocolbuffers/protobuf/issues/2283
+Additionally rules_proto_grpc 4.5.0 (and all pre 5.0.0/bazel mod versions) do not support pyi generation.
+To supply debug symbols an additional genrule is needed to compile them manually with protoc.
 """
 
-py_library(
-    name = "opensearch_protos_schemas",
+genrule(
+    name = "generate_pyi_files",
     srcs = [
-        "//protos/schemas:common_python_proto",
-        "//protos/schemas:document_python_proto",
-        "//protos/schemas:search_python_proto",
+        "//protos/schemas:common.proto",
+        "//protos/schemas:document.proto",
+        "//protos/schemas:search.proto",
+        "//protos/services:document_service.proto",
+        "//protos/services:search_service.proto",
+        "@com_google_protobuf//:well_known_type_protos",
     ],
-)
-
-py_library(
-    name = "opensearch_protos_services",
-    srcs = [
-        "//protos/services:document_service_python_proto",
-        "//protos/services:search_service_python_proto",
+    outs = [
+        "python_pyi/schemas/common_pb2.pyi",
+        "python_pyi/schemas/document_pb2.pyi", 
+        "python_pyi/schemas/search_pb2.pyi",
+        "python_pyi/services/document_service_pb2.pyi",
+        "python_pyi/services/search_service_pb2.pyi",
     ],
-    deps = [":opensearch_protos_schemas"],
+    cmd = """        
+        # Use protoc with --pyi_out to generate debug symbols
+        $(location @com_google_protobuf//:protoc) \
+            --proto_path=. \
+            --proto_path=external/com_google_protobuf/src \
+            --pyi_out=$(RULEDIR) \
+            protos/schemas/common.proto \
+            protos/schemas/document.proto \
+            protos/schemas/search.proto \
+            protos/services/document_service.proto \
+            protos/services/search_service.proto
+
+        mkdir -p $(RULEDIR)/python_pyi/
+        mv $(RULEDIR)/protos/schemas $(RULEDIR)/python_pyi
+        mv $(RULEDIR)/protos/services $(RULEDIR)/python_pyi
+    """,
+    tools = ["@com_google_protobuf//:protoc"],
 )
 
 genrule(
-    name = "fix_schemas_imports",
-    srcs = [":opensearch_protos_schemas"],
+    name = "fix_python_imports",
+    srcs = [":python_schemas", ":python_services", ":generate_pyi_files"],
     outs = [
         "opensearch/protobufs/schemas/common_pb2.py",
+        "opensearch/protobufs/schemas/common_pb2.pyi",
         "opensearch/protobufs/schemas/document_pb2.py",
+        "opensearch/protobufs/schemas/document_pb2.pyi",
         "opensearch/protobufs/schemas/search_pb2.py",
-        "opensearch/protobufs/schemas/__init__.py",
-    ],
-    cmd = """
-        mkdir -p $(RULEDIR)/opensearch/protobufs/schemas
-
-        for pb in common document search; do
-            cp $(BINDIR)/protos/schemas/$${pb}_python_proto_pb/protos/schemas/$${pb}_pb2.py $(RULEDIR)/opensearch/protobufs/schemas/
-            sed -i \
-                -e 's/from protos\\.schemas import \\([^_ ]\\+\\)_pb2 as protos_dot_schemas_dot_\\1__pb2/from opensearch.protobufs.schemas import \\1_pb2/g' \
-                $(RULEDIR)/opensearch/protobufs/schemas/$${pb}_pb2.py
-        done
-
-        touch $(RULEDIR)/opensearch/protobufs/schemas/__init__.py
-    """,
-)
-
-genrule(
-    name = "fix_services_imports",
-    srcs = [":opensearch_protos_services"],
-    outs = [
+        "opensearch/protobufs/schemas/search_pb2.pyi",
         "opensearch/protobufs/services/document_service_pb2.py",
+        "opensearch/protobufs/services/document_service_pb2.pyi",
         "opensearch/protobufs/services/document_service_pb2_grpc.py",
         "opensearch/protobufs/services/search_service_pb2.py",
+        "opensearch/protobufs/services/search_service_pb2.pyi",
         "opensearch/protobufs/services/search_service_pb2_grpc.py",
+        "opensearch/protobufs/schemas/__init__.py",
         "opensearch/protobufs/services/__init__.py",
     ],
     cmd = """
-        mkdir -p $(RULEDIR)/opensearch/protobufs/services
+        mkdir -p $(RULEDIR)/opensearch/protobufs/
+        
+        cp -rL $(BINDIR)/python_schemas_pb/protos/schemas $(RULEDIR)/opensearch/protobufs
+        cp -rL $(BINDIR)/python_services_pb/protos/services $(RULEDIR)/opensearch/protobufs
 
-        for pb in document_service search_service; do
-            cp $(BINDIR)/protos/services/$${pb}_python_proto_pb/protos/services/$${pb}_pb2.py $(RULEDIR)/opensearch/protobufs/services/
-            cp $(BINDIR)/protos/services/$${pb}_python_proto_pb/protos/services/$${pb}_pb2_grpc.py $(RULEDIR)/opensearch/protobufs/services/
-
-            # Fix imports in both pb2 and pb2_grpc files
-            sed -i \
-                -e 's/from protos\\.schemas import/from opensearch.protobufs.schemas import/g' \
-                $(RULEDIR)/opensearch/protobufs/services/$${pb}_pb2*.py
-        done
+        cp -rL $(BINDIR)/python_pyi/schemas $(RULEDIR)/opensearch/protobufs
+        cp -rL $(BINDIR)/python_pyi/services $(RULEDIR)/opensearch/protobufs
 
+        # Find and replace "from protos" with "from opensearch.protobufs" in all files
+        find $(RULEDIR)/opensearch/protobufs/ -type f -exec sed -i 's/from protos/from opensearch.protobufs/g' {} +
+        
+        touch $(RULEDIR)/opensearch/protobufs/schemas/__init__.py
         touch $(RULEDIR)/opensearch/protobufs/services/__init__.py
     """,
 )
 
-py_library(
-    name = "fixed_schemas",
-    srcs = [":fix_schemas_imports"],
-    imports = ["."],
-    visibility = ["//visibility:public"],
-)
-
-py_library(
-    name = "fixed_services",
-    srcs = [":fix_services_imports"],
-    imports = ["."],
-    deps = [":fixed_schemas"],
-    visibility = ["//visibility:public"],
-)
-
 py_wheel(
     name = "opensearch_protos_wheel",
     distribution = "opensearch-protobufs",
@@ -159,7 +157,6 @@ py_wheel(
     platform = "any",
     python_requires = ">=3.10",
     deps = [
-        ":fixed_schemas",
-        ":fixed_services",
+        ":fix_python_imports",
     ],
 )

CHANGELOG.md

@@ -4,6 +4,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ## [Unreleased]
 ### Added
+- Add .pyi type hints to the python wheel ([#253](https://github.com/opensearch-project/opensearch-protobufs/pull/253))
 
 ### Changed
 

Dockerfile

@@ -90,7 +90,7 @@ RUN ./gradlew :plugins:transport-grpc:internalClusterTest -Drepos.mavenLocal
 
 FROM build-bazel AS build-bazel-python
 
-RUN bazel build //:python_protos_all
+RUN bazel build //:python_schemas //:python_services
 
 FROM build-bazel-python AS package-bazel-python
 

README.md

@@ -17,7 +17,7 @@ Generate protobuf libraries for your preferred language:
 bazel build //:java_protos_all
 
 # Python
-bazel build //:python_protos_all
+bazel build //:python_schemas //:python_services
 
 # Go
 bazel build //:go_protos_all

protos/schemas/BUILD.bazel

@@ -1,7 +1,6 @@
 # protos/schemas/BUILD.bazel
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@rules_proto_grpc//java:defs.bzl", "java_proto_library")
-load("@rules_proto_grpc//python:defs.bzl", "python_proto_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
 
 package(default_visibility = ["//visibility:public"])
@@ -47,21 +46,6 @@ java_proto_library(
     deps = [":common_java_proto"],
 )
 
-python_proto_library(
-    name = "common_python_proto",
-    protos = [":common_proto"],
-)
-
-python_proto_library(
-    name = "document_python_proto",
-    protos = [":document_proto"],
-)
-
-python_proto_library(
-    name = "search_python_proto",
-    protos = [":search_proto"],
-)
-
 go_proto_library(
     name = "common_go_proto",
     protos = [":common_proto"],

protos/services/BUILD.bazel

@@ -1,7 +1,6 @@
 # protos/services/BUILD.bazel
 load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@rules_proto_grpc//java:defs.bzl", "java_grpc_library")
-load("@rules_proto_grpc//python:defs.bzl", "python_grpc_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_grpc_library")
 
 package(default_visibility = ["//visibility:public"])
@@ -30,16 +29,6 @@ java_grpc_library(
     deps = ["//protos/schemas:search_java_proto"],
 )
 
-python_grpc_library(
-    name = "document_service_python_proto",
-    protos = [":document_service_proto"],
-)
-
-python_grpc_library(
-    name = "search_service_python_proto",
-    protos = [":search_service_proto"],
-)
-
 go_grpc_library(
     name = "document_service_go_proto",
     protos = [":document_service_proto"],