From 711885db8ffd4ff201a5ff7f00c431f36e6ccdf4 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <jhaberman@gmail.com>
Date: Sat, 26 Feb 2022 09:48:05 -0800
Subject: [PATCH] Added py_extension().

---
 bazel/py_extension.bzl     | 39 ++++++++++++++++++++++++++++++++++++++
 bazel/py_proto_library.bzl |  2 +-
 bazel/workspace_defs.bzl   | 18 +++++++++---------
 bazel/workspace_deps.bzl   |  4 ++--
 python/BUILD               | 28 +++++++--------------------
 5 files changed, 58 insertions(+), 33 deletions(-)
 create mode 100644 bazel/py_extension.bzl

diff --git a/bazel/py_extension.bzl b/bazel/py_extension.bzl
new file mode 100644
index 0000000000..7c1adf4146
--- /dev/null
+++ b/bazel/py_extension.bzl
@@ -0,0 +1,39 @@
+
+load(
+    "//bazel:build_defs.bzl",
+    "UPB_DEFAULT_COPTS",
+)
+
+def py_extension(name, srcs, deps):
+    version_script = name + "_version_script.lds"
+    symbol = "PyInit_" + name
+    native.genrule(
+        name = "gen_" + version_script,
+        outs = [version_script],
+        cmd = "echo 'message { global: " + symbol + "; local: *; };' > $@",
+    )
+
+    native.cc_binary(
+        name = name,
+        srcs = srcs,
+        copts = UPB_DEFAULT_COPTS + [
+            # The Python API requires patterns that are ISO C incompatible, like
+            # casts between function pointers and object pointers.
+            "-Wno-pedantic",
+        ],
+        # We use a linker script to hide all symbols except the entry point for
+        # the module.
+        linkopts = select({
+            "@platforms//os:linux": ["-Wl,--version-script,$(location :" + version_script + ")"],
+            "@platforms//os:macos": [
+                "-Wl,-exported_symbol",
+                "-Wl,_" + symbol,
+            ],
+        }),
+        linkshared = True,
+        linkstatic = True,
+        deps = deps + [
+            ":" + version_script,
+            "@system_python//:python_headers",
+        ],
+    )
diff --git a/bazel/py_proto_library.bzl b/bazel/py_proto_library.bzl
index 9f99112346..8ef4d9708d 100644
--- a/bazel/py_proto_library.bzl
+++ b/bazel/py_proto_library.bzl
@@ -103,7 +103,7 @@ def _py_proto_library_aspect_impl(target, ctx):
     )
     outs_depset = depset(srcs)
     return [
-        PyInfo(transitive_sources = outs_depset)
+        PyInfo(transitive_sources = outs_depset),
     ]
 
 _py_proto_library_aspect = aspect(
diff --git a/bazel/workspace_defs.bzl b/bazel/workspace_defs.bzl
index 2f596f01fb..03655b2409 100644
--- a/bazel/workspace_defs.bzl
+++ b/bazel/workspace_defs.bzl
@@ -53,17 +53,17 @@ toolchain(
 """
 
 def _get_config_var(repository_ctx, name):
-  py_program = "import sysconfig; print(sysconfig.get_config_var('%s'), end='')"
-  result = repository_ctx.execute(["python3", "-c", py_program % (name)])
-  if result.return_code != 0:
-    fail("No python3 executable available on the system")
-  return result.stdout
+    py_program = "import sysconfig; print(sysconfig.get_config_var('%s'), end='')"
+    result = repository_ctx.execute(["python3", "-c", py_program % (name)])
+    if result.return_code != 0:
+        fail("No python3 executable available on the system")
+    return result.stdout
 
 def _python_headers_impl(repository_ctx):
-  path = _get_config_var(repository_ctx, "INCLUDEPY")
-  repository_ctx.symlink(path, "python")
-  python3 = repository_ctx.which("python3")
-  repository_ctx.file("BUILD.bazel", _build_file % python3)
+    path = _get_config_var(repository_ctx, "INCLUDEPY")
+    repository_ctx.symlink(path, "python")
+    python3 = repository_ctx.which("python3")
+    repository_ctx.file("BUILD.bazel", _build_file % python3)
 
 # The system_python() repository rule exposes Python headers from the system.
 #
diff --git a/bazel/workspace_deps.bzl b/bazel/workspace_deps.bzl
index a5a3564449..db3afd7a79 100644
--- a/bazel/workspace_deps.bzl
+++ b/bazel/workspace_deps.bzl
@@ -23,10 +23,10 @@ def upb_deps():
             "rm python/google/protobuf/__init__.py",
             "rm python/google/protobuf/pyext/__init__.py",
             "rm python/google/protobuf/internal/__init__.py",
-        ]
+        ],
     )
 
-    rules_python_version = "740825b7f74930c62f44af95c9a4c1bd428d2c53" # Latest @ 2021-06-23
+    rules_python_version = "740825b7f74930c62f44af95c9a4c1bd428d2c53"  # Latest @ 2021-06-23
 
     maybe(
         http_archive,
diff --git a/python/BUILD b/python/BUILD
index 8126f52200..687ebe6003 100644
--- a/python/BUILD
+++ b/python/BUILD
@@ -31,6 +31,10 @@ load(
     "//bazel:py_proto_library.bzl",
     "py_proto_library",
 )
+load(
+    "//bazel:py_extension.bzl",
+    "py_extension",
+)
 load(
     "@rules_python//python:packaging.bzl",
     "py_wheel",
@@ -38,8 +42,8 @@ load(
 
 licenses(["notice"])
 
-cc_binary(
-    name = "message",
+py_extension(
+    name = "_message",
     srcs = [
         "convert.c",
         "convert.h",
@@ -61,24 +65,7 @@ cc_binary(
         "repeated.c",
         "repeated.h",
     ],
-    copts = UPB_DEFAULT_COPTS + [
-        # The Python API requires patterns that are ISO C incompatible, like
-        # casts between function pointers and object pointers.
-        "-Wno-pedantic",
-    ],
-    # We use a linker script to hide all symbols except the entry point for
-    # the module.
-    linkopts = select({
-        "@platforms//os:linux": ["-Wl,--version-script,$(location :version_script.lds)"],
-        "@platforms//os:macos": [
-            "-Wl,-exported_symbol",
-            "-Wl,_PyInit__message",
-        ],
-    }),
-    linkshared = True,
-    linkstatic = True,
     deps = [
-        ":version_script.lds",
         "//:descriptor_upb_proto_reflection",
         "//:reflection",
         "//:textformat",
@@ -86,7 +73,6 @@ cc_binary(
         "//upb/util:compare",
         "//upb/util:def_to_proto",
         "//upb/util:required_fields",
-        "@system_python//:python_headers",
     ],
 )
 
@@ -108,7 +94,7 @@ EXT_SUFFIX = ".abi3.so"
 
 genrule(
     name = "copy_message",
-    srcs = [":message"],
+    srcs = [":_message"],
     outs = ["google/protobuf/pyext/_message" + EXT_SUFFIX],
     cmd = "cp $< $@",
 )