From 22f89e02d10225c47559f2082e368a5b18c9ed95 Mon Sep 17 00:00:00 2001
From: Tim Ohliger <ohliger.tim@gmail.com>
Date: Thu, 15 Aug 2024 10:51:53 +0200
Subject: [PATCH] Moved namespace fix to autodoc event callbacks to fix imports
 in stubs.

---
 cpp/pybind/docstring.cpp | 12 ++----------
 docs/conf.py             | 23 +++++++++++++++++++++++
 2 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/cpp/pybind/docstring.cpp b/cpp/pybind/docstring.cpp
index fce26d58688..e9d03711d35 100644
--- a/cpp/pybind/docstring.cpp
+++ b/cpp/pybind/docstring.cpp
@@ -289,17 +289,9 @@ std::string FunctionDoc::ToGoogleDocString() const {
     return rc.str();
 }
 
-std::string FunctionDoc::NamespaceFix(const std::string& s) {
-    std::string rc = std::regex_replace(s, std::regex("::(\\S)"), ".$1");
-    rc = std::regex_replace(rc, std::regex("open3d\\.(cpu|cuda)\\.pybind\\."),
-                            "open3d.");
-    return rc;
-}
-
 std::string FunctionDoc::StringCleanAll(std::string& s,
                                         const std::string& white_space) {
     std::string rc = utility::StripString(s, white_space);
-    rc = NamespaceFix(rc);
     return rc;
 }
 
@@ -313,7 +305,7 @@ ArgumentDoc FunctionDoc::ParseArgumentToken(const std::string& argument_token) {
     std::smatch matches;
     if (std::regex_search(argument_token, matches, rgx_with_default)) {
         argument_doc.name_ = matches[1].str();
-        argument_doc.type_ = NamespaceFix(matches[2].str());
+        argument_doc.type_ = matches[2].str();
         argument_doc.default_ = matches[3].str();
 
         // Handle long default value. Long default has multiple lines and thus
@@ -335,7 +327,7 @@ ArgumentDoc FunctionDoc::ParseArgumentToken(const std::string& argument_token) {
                 "([A-Za-z_][A-Za-z\\d_:\\.\\[\\]\\(\\) ,]*)");
         if (std::regex_search(argument_token, matches, rgx_without_default)) {
             argument_doc.name_ = matches[1].str();
-            argument_doc.type_ = NamespaceFix(matches[2].str());
+            argument_doc.type_ = matches[2].str();
         }
     }
 
diff --git a/docs/conf.py b/docs/conf.py
index eac752217c7..61d9234ec72 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -264,8 +264,31 @@ def skip(app, what, name, obj, would_skip, options):
     return would_skip
 
 
+def fix_namespace(text):
+    """Fixes namespace in a string by removing .[cpu|cuda].pybind."""
+    return (text.replace("open3d.cpu.pybind.", "open3d.").replace(
+        "open3d.cuda.pybind.cuda.", "open3d.") if text is not None else None)
+
+
+def process_signature(app, what, name, obj, options, signature,
+                      return_annotation):
+    """Fixes namespace in signature by removing .[cpu|cuda].pybind."""
+    return (
+        fix_namespace(signature),
+        fix_namespace(return_annotation),
+    )
+
+
+def process_docstring(app, what, name, obj, options, lines):
+    """Fixes namespace in docstring by removing .[cpu|cuda].pybind."""
+    for i, line in enumerate(lines):
+        lines[i] = fix_namespace(line)
+
+
 def setup(app):
     app.connect("autodoc-skip-member", skip)
+    app.connect("autodoc-process-signature", process_signature)
+    app.connect("autodoc-process-docstring", process_docstring)
     # Add Google analytics
     app.add_js_file("https://www.googletagmanager.com/gtag/js?id=G-3TQPKGV6Z3",
                     **{'async': 'async'})