From fbcd42b9f00afef3ecd1dd56d91a2bb11ad9e350 Mon Sep 17 00:00:00 2001
From: Joshua Haberman <haberman@google.com>
Date: Fri, 29 Apr 2022 10:07:05 -0700
Subject: [PATCH] Added a test that the Python module is importable, and fixed
 several issues that were required to make the test pass:

1. For some reason the version script was not working, it was failing to export the main symbol for the Python extension.  I fixed this by using the `visibility` attribute instead to export the `PyInit__message` function.
2. We were not properly stripping the `python/dist/` prefix for the C module, which was making the module exported under the name `dist.google._upb` instead of `google._upb`.
3. The `py_library()` rule was failing to actually propagate the module file into the wheel, so I just removed it.

PiperOrigin-RevId: 445446611
---
 .github/workflows/python_tests.yml |  2 ++
 python/dist/BUILD.bazel            |  5 ++++-
 python/dist/dist.bzl               | 10 +---------
 python/protobuf.c                  |  2 +-
 python/py_extension.bzl            | 14 ++------------
 5 files changed, 10 insertions(+), 23 deletions(-)

diff --git a/.github/workflows/python_tests.yml b/.github/workflows/python_tests.yml
index a277f80fad..4c5098b18f 100644
--- a/.github/workflows/python_tests.yml
+++ b/.github/workflows/python_tests.yml
@@ -70,3 +70,5 @@ jobs:
           echo "VIRTUAL ENV:" $VIRTUAL_ENV
       - name: Install Wheels
         run: pip install -vvv --no-index --find-links wheels protobuf
+      - name: Test that module is importable
+        run: python -c 'from google._upb import _message; assert "google._upb._message.MessageMeta" in str(_message.MessageMeta)'
diff --git a/python/dist/BUILD.bazel b/python/dist/BUILD.bazel
index c8a623502f..9ee1f29ebb 100644
--- a/python/dist/BUILD.bazel
+++ b/python/dist/BUILD.bazel
@@ -104,7 +104,10 @@ py_wheel(
         "//python:limited_api_3.10": "cp310",
         "//conditions:default": "system",
     }),
-    strip_path_prefixes = ["python/"],
+    strip_path_prefixes = [
+        "python/dist/",
+        "python/",
+    ],
     version = PROTOBUF_VERSION,
     deps = [
         ":message_mod",
diff --git a/python/dist/dist.bzl b/python/dist/dist.bzl
index 6760a381b3..fc725f3387 100644
--- a/python/dist/dist.bzl
+++ b/python/dist/dist.bzl
@@ -76,20 +76,12 @@ _py_dist_module_rule = rule(
 )
 
 def py_dist_module(name, module_name, extension):
-    file_rule = name + "_file"
     _py_dist_module_rule(
-        name = file_rule,
+        name = name,
         module_name = module_name,
         extension = extension,
     )
 
-    # TODO(haberman): needed?
-    native.py_library(
-        name = name,
-        data = [":" + file_rule],
-        imports = ["."],
-    )
-
 def _py_dist_transition_impl(settings, attr):
     _ignore = (settings)  # @unused
     transitions = []
diff --git a/python/protobuf.c b/python/protobuf.c
index e41657994a..53bfe3a8a7 100644
--- a/python/protobuf.c
+++ b/python/protobuf.c
@@ -322,7 +322,7 @@ PyObject* PyUpb_Forbidden_New(PyObject* cls, PyObject* args, PyObject* kwds) {
 // Module Entry Point
 // -----------------------------------------------------------------------------
 
-PyMODINIT_FUNC PyInit__message(void) {
+__attribute__((visibility("default"))) PyMODINIT_FUNC PyInit__message(void) {
   PyObject* m = PyModule_Create(&module_def);
   if (!m) return NULL;
 
diff --git a/python/py_extension.bzl b/python/py_extension.bzl
index 77a50ffe27..2b43e43ffd 100644
--- a/python/py_extension.bzl
+++ b/python/py_extension.bzl
@@ -11,27 +11,17 @@ def py_extension(name, srcs, copts, deps = []):
       deps: Libraries that the target depends on
     """
 
-    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 + "_binary",
         srcs = srcs,
-        copts = copts,
+        copts = copts + ["-fvisibility=hidden"],
         linkopts = select({
             "//python/dist:osx-x86_64_cpu": ["-undefined", "dynamic_lookup"],
             "//conditions:default": [],
         }),
         linkshared = True,
         linkstatic = True,
-        deps = deps + [
-            ":" + version_script,
-        ] + select({
+        deps = deps + select({
             "//python:limited_api_3.7": ["@python-3.7.0//:python_headers"],
             "//python:full_api_3.7_win32": ["@nuget_python_i686_3.7.0//:python_full_api"],
             "//python:full_api_3.7_win64": ["@nuget_python_x86-64_3.7.0//:python_full_api"],