feat: add support for submodules (including recursive submodules) in pybind_extensions

vam-google · vam-google · commit f7b469ea0b57 · 2025-04-05T22:26:45.000-07:00
Also changes approach in import shim from regexporting all non-private symbols under new namespace to direct inserting of the original modules under proper names in runtime.
Also add some extra testing, such as relative import testing.
diff --git a/pybind/BUILD b/pybind/BUILD
@@ -132,9 +132,6 @@ cc_test(
 pybind_extension(
     name = "pybind",
     srcs = ["pybind.cc"],
-    additional_exported_symbols = [
-        "_EXTRA_SYMBOL",
-    ],
     common_lib_packages = [
         "",
         "pywrap_external/python",  # non-existent package to test that its Ok
@@ -154,9 +151,6 @@ pybind_extension(
 pybind_extension(
     name = "pybind_copy",
     srcs = ["pybind_copy.cc"],
-    additional_exported_symbols = [
-        "_EXTRA_SYMBOL",
-    ],
     common_lib_packages = [
         "",
         "pywrap_external",
@@ -253,7 +247,7 @@ pywrap_library(
             "//conditions:default": ":linux_framework.lds",
         }),
     },
-    pywrap_count = 3,
+    pywrap_count = 4,
     pywrap_lib_filter = ":third_library",
     starlark_only_pywrap_count = 1,
     visibility = ["//visibility:public"],
@@ -263,6 +257,7 @@ pywrap_library(
         ":py_library_3",
         ":py_library_4",
         ":pybind_with_starlark_only",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
 )
 
@@ -291,7 +286,9 @@ py_test(
         # not needed, since duplicated, but passed to check that duplication is
         # ok
         "py_library_1",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
+#    legacy_create_init = False,
 )
 
 py_test(
diff --git a/pybind/pybind.cc b/pybind/pybind.cc
@@ -4,6 +4,10 @@
 #include "pybind11/pybind11.h"
 #include "fifth_library.h"
 
+int sub_sub_func(int x) {
+  return x << 1;
+}
+
 PYBIND11_MODULE(pybind, m) {
   fifth_func(); // needed to  make sure the dynamic library is loaded on Windows
   m.doc() = "pybind11 example plugin";
@@ -12,4 +16,10 @@ PYBIND11_MODULE(pybind, m) {
   m.def("third_func", &third_func, "The third function");
   m.def("second_global_func", &second_global_func, "The second_global function");
   m.attr("_EXTRA_SYMBOL") = pybind11::int_(123);
+
+  auto subM = m.def_submodule("sub", "submodule");
+  subM.def("second_func", &second_func, "");
+  auto subSubM = subM.def_submodule("sub_sub", "sub submodule");
+  subSubM.def("sub_sub_func", &sub_sub_func, "");
+  subSubM.def("_sub_sub_private_func", &sub_sub_func, "private sub func");
 }
diff --git a/pybind/pybind_py_test.py b/pybind/pybind_py_test.py
@@ -6,7 +6,12 @@
 import pybind.pybind_copy
 from pybind.pybind import _EXTRA_SYMBOL as REGULAR_EXTRA_SYMBOL
 from pybind.pybind_copy import _EXTRA_SYMBOL as REGULAR_COPY_EXTRA_SYMBOL
-
+from pybind.pybind.sub import second_func as sub_second_func
+import pybind.pybind.sub
+from pybind.pybind.sub.sub_sub import *
+from pybind.pybind.sub.sub_sub import _sub_sub_private_func
+from pybind.sub_pybind.relative_import_lib import call_nested_pyind_func
+from pybind.sub_pybind.relative_import_lib import sub_sub_private_func
 
 class PybindTest(unittest.TestCase):
   def _read_file(self, filename, mode="r"):
@@ -48,6 +53,15 @@ def test_pybind_first(self):
     self.assertEqual(self._read_file("pybind/static_resource.txt"),
                      "A static resource file under pybind dir")
 
+    print("14: Submodules")
+    self.assertEqual(sub_second_func(1), 5)
+    self.assertEqual(pybind.pybind.sub.second_func(1), 6)
+    self.assertEqual(pybind.pybind.sub.sub_sub.sub_sub_func(3), 6)
+    self.assertEqual(_sub_sub_private_func(5), 10)
+
+    print("15: Nested pybinds and relative imports")
+    self.assertEqual(call_nested_pyind_func(6), 3)
+    self.assertEqual(sub_sub_private_func(5), 10)
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/pybind/sub_pybind/BUILD b/pybind/sub_pybind/BUILD
@@ -0,0 +1,26 @@
+load("//rules_pywrap:pywrap.bzl", "pybind_extension")
+
+pybind_extension(
+    name = "nested_pybind",
+    srcs = ["nested_pybind.cc"],
+    common_lib_packages = [
+        "",
+        "pywrap_external/python",  # non-existent package to test that its Ok
+        "pybind",
+        "pybind_external",
+        "pywrap_external",
+    ],
+#    visibility = ["//visibility:public"],
+    deps = [
+        "//pybind:fifth_library",
+    ],
+)
+
+py_library(
+    name = "relative_import_lib",
+    srcs = ["relative_import_lib.py"],
+    deps = [
+        ":nested_pybind",
+    ],
+    visibility = ["//visibility:public"]
+)
diff --git a/pybind/sub_pybind/nested_pybind.cc b/pybind/sub_pybind/nested_pybind.cc
@@ -0,0 +1,11 @@
+#include "pybind11/pybind11.h"
+#include "pybind/fifth_library.h"
+
+int nested_pybind_func(int x) {
+  return x >> 1;
+}
+
+PYBIND11_MODULE(nested_pybind, m) {
+  fifth_func(); // needed to  make sure the dynamic library is loaded on Windows
+  m.def("nested_pybind_func", &nested_pybind_func, "");
+}
diff --git a/pybind/sub_pybind/relative_import_lib.py b/pybind/sub_pybind/relative_import_lib.py
@@ -0,0 +1,8 @@
+from .nested_pybind import nested_pybind_func
+from ..pybind.sub.sub_sub import _sub_sub_private_func
+
+def call_nested_pyind_func(x):
+  return nested_pybind_func(x)
+
+def sub_sub_private_func(x):
+  return _sub_sub_private_func(x)
diff --git a/pybind_external/BUILD b/pybind_external/BUILD
@@ -8,6 +8,7 @@ pywrap_library(
         "//pybind",
         "//pybind:pybind_cc_only",
         "//pybind:pybind_copy",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
 )
 
@@ -22,6 +23,7 @@ py_test(
     visibility = ["//:__pkg__"],
     deps = [
         ":pybind_external_aggregated",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
 )
 
@@ -32,6 +34,7 @@ py_test(
     visibility = ["//:__pkg__"],
     deps = [
         "//pywrap_external:pywrap_external_aggregated",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
 )
 
diff --git a/pybind_external/pybind_external_py_test.py b/pybind_external/pybind_external_py_test.py
@@ -6,7 +6,12 @@
 import pybind.pybind_copy
 from pybind.pybind import _EXTRA_SYMBOL as REGULAR_EXTRA_SYMBOL
 from pybind.pybind_copy import _EXTRA_SYMBOL as REGULAR_COPY_EXTRA_SYMBOL
-
+from pybind.pybind.sub import second_func as sub_second_func
+import pybind.pybind.sub
+from pybind.pybind.sub.sub_sub import *
+from pybind.pybind.sub.sub_sub import _sub_sub_private_func
+from pybind.sub_pybind.relative_import_lib import call_nested_pyind_func
+from pybind.sub_pybind.relative_import_lib import sub_sub_private_func
 
 class PybindTest(unittest.TestCase):
   def _read_file(self, filename, mode="r"):
@@ -48,6 +53,15 @@ def test_pybind_first(self):
     self.assertEqual(self._read_file("pybind/static_resource.txt"),
                      "A static resource file under pybind dir")
 
+    print("14: Submodules")
+    self.assertEqual(sub_second_func(1), 5)
+    self.assertEqual(pybind.pybind.sub.second_func(1), 6)
+    self.assertEqual(pybind.pybind.sub.sub_sub.sub_sub_func(3), 6)
+    self.assertEqual(_sub_sub_private_func(5), 10)
+
+    print("15: Nested pybinds and relative imports")
+    self.assertEqual(call_nested_pyind_func(6), 3)
+    self.assertEqual(sub_sub_private_func(5), 10)
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/pybind_external/pywrap_external_aggregated_binaries_test.py b/pybind_external/pywrap_external_aggregated_binaries_test.py
@@ -31,6 +31,9 @@ def test_pywrap_binaries(self):
         (
             "/pywrap_external/{lib}pywrap_external_aggregated__starlark_only_common.{extension}",
             ""),
+        ('/pybind/sub_pybind/nested_pybind.py', ''),
+        ('/pywrap_external/nested_pybind.{pyextension}',
+         '/pybind/sub_pybind/nested_pybind.{pyextension}'),
     ]
 
     pyextension = "so"
@@ -40,17 +43,22 @@ def test_pywrap_binaries(self):
     if "Windows" in system:
       relative_wheel_locations.extend([
           ("/pywrap_external/framework.2.dll", "/pybind/framework.2.dll"),
-          ("/pywrap_external/framework.2.dll.if.lib", "/pybind/framework.2.dll.if.lib"),
-          ("/pywrap_external/pywrap_external_aggregated__starlark_only_common.dll.if.lib", ""),
-          ("/pywrap_external/pywrap_external_aggregated_common.dll.if.lib", "/pywrap_external/pywrap_external_aggregated_common.dll.if.lib"),
+          ("/pywrap_external/framework.2.dll.if.lib",
+           "/pybind/framework.2.dll.if.lib"),
+          (
+              "/pywrap_external/pywrap_external_aggregated__starlark_only_common.dll.if.lib",
+              ""),
+          ("/pywrap_external/pywrap_external_aggregated_common.dll.if.lib",
+           "/pywrap_external/pywrap_external_aggregated_common.dll.if.lib"),
       ])
       pyextension = "pyd"
       extension = "dll"
       lib_prefix = ""
     elif "Darwin" in system:
       extension = "dylib"
       relative_wheel_locations.extend([
-          ("/pywrap_external/libframework.2.dylib", "/pybind/libframework.2.dylib"),
+          ("/pywrap_external/libframework.2.dylib",
+           "/pybind/libframework.2.dylib"),
       ])
     else:
       relative_wheel_locations.extend([
diff --git a/pywrap_external/BUILD b/pywrap_external/BUILD
@@ -17,6 +17,7 @@ pywrap_library(
         "//pybind",
         "//pybind:pybind_cc_only",
         "//pybind:pybind_copy",
+        "//pybind/sub_pybind:relative_import_lib",
     ],
 )
 
diff --git a/rules_pywrap/pybind_extension.py.tpl b/rules_pywrap/pybind_extension.py.tpl
@@ -1,13 +1,18 @@
-def __update_globals(pywrap_m):
+from sys import modules
+from types import ModuleType
+
+def __update_globals(new_import_path, pywrap_m):
   if hasattr(pywrap_m, '__all__'):
     all_names = pywrap_m.__all__
   else:
     all_names = [name for name in dir(pywrap_m) if not name.startswith('_')]
 
-  extra_names = []  # template_val
-  all_names.extend(extra_names)
-  globals().update({name: getattr(pywrap_m, name) for name in all_names})
-
+  modules[new_import_path] = pywrap_m
+  for name in all_names:
+    sub_pywrap = getattr(pywrap_m, name)
+    if isinstance(sub_pywrap, ModuleType):
+      sub_name = sub_pywrap.__name__[len(pywrap_m.__name__):]
+      __update_globals(new_import_path + sub_name, sub_pywrap)
 
 def __try_import():
   imports_paths = []  # template_val
@@ -16,7 +21,7 @@ def __try_import():
   for import_path in imports_paths:
     try:
       pywrap_m = __import__(import_path, fromlist=["*"])
-      __update_globals(pywrap_m)
+      __update_globals(__name__, pywrap_m)
       return
     except ImportError as e:
       exceptions.append(str(e))
diff --git a/rules_pywrap/pywrap.bzl b/rules_pywrap/pywrap.bzl
@@ -752,6 +752,8 @@ def pybind_extension(
         linkopts = [],
         starlark_only = False,
         **kwargs):
+    # For backward compatibility that I don't want to mess with
+    _ignore = [additional_exported_symbols]
     cc_library_name = "_%s_cc_library" % name
     native.cc_library(
         name = cc_library_name,
@@ -793,7 +795,6 @@ def pybind_extension(
             name = name,
             deps = ["%s" % cc_library_name],
             common_lib_packages = common_lib_packages,
-            additional_exported_symbols = additional_exported_symbols,
             starlark_only = starlark_only,
             testonly = testonly,
             compatible_with = compatible_with,
@@ -807,9 +808,6 @@ def _pywrap_info_wrapper_impl(ctx):
 
     py_stub = ctx.actions.declare_file("%s.py" % ctx.attr.name)
     substitutions = {}
-
-    additional_exported_symbols = ctx.attr.additional_exported_symbols
-
     py_pkgs = []
     for pkg in ctx.attr.common_lib_packages:
         if pkg:
@@ -819,10 +817,6 @@ def _pywrap_info_wrapper_impl(ctx):
         val = "imports_paths = %s # template_val" % py_pkgs
         substitutions["imports_paths = []  # template_val"] = val
 
-    if additional_exported_symbols:
-        val = "extra_names = %s # template_val" % additional_exported_symbols
-        substitutions["extra_names = []  # template_val"] = val
-
     ctx.actions.expand_template(
         template = ctx.file.py_stub_src,
         output = py_stub,
@@ -854,10 +848,6 @@ _pywrap_info_wrapper = rule(
             allow_single_file = True,
             default = Label("//rules_pywrap:pybind_extension.py.tpl"),
         ),
-        "additional_exported_symbols": attr.string_list(
-            mandatory = False,
-            default = [],
-        ),
         "starlark_only": attr.bool(mandatory = False, default = False),
     },
     implementation = _pywrap_info_wrapper_impl,

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@ pywrap_library(`
`8`	`8`	`"//pybind",`
`9`	`9`	`"//pybind:pybind_cc_only",`
`10`	`10`	`"//pybind:pybind_copy",`
	`11`	`+ "//pybind/sub_pybind:relative_import_lib",`
`11`	`12`	`],`
`12`	`13`	`)`
`13`	`14`
`@@ -22,6 +23,7 @@ py_test(`
`22`	`23`	`visibility = ["//:__pkg__"],`
`23`	`24`	`deps = [`
`24`	`25`	`":pybind_external_aggregated",`
	`26`	`+ "//pybind/sub_pybind:relative_import_lib",`
`25`	`27`	`],`
`26`	`28`	`)`
`27`	`29`
`@@ -32,6 +34,7 @@ py_test(`
`32`	`34`	`visibility = ["//:__pkg__"],`
`33`	`35`	`deps = [`
`34`	`36`	`"//pywrap_external:pywrap_external_aggregated",`
	`37`	`+ "//pybind/sub_pybind:relative_import_lib",`
`35`	`38`	`],`
`36`	`39`	`)`
`37`	`40`
Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@ pywrap_library(`
`17`	`17`	`"//pybind",`
`18`	`18`	`"//pybind:pybind_cc_only",`
`19`	`19`	`"//pybind:pybind_copy",`
	`20`	`+ "//pybind/sub_pybind:relative_import_lib",`
`20`	`21`	`],`
`21`	`22`	`)`
`22`	`23`