Added type annotations to func signature output by ToGoogleDocString.

isl-org · Aug 14, 2024 · 703bead · 703bead
1 parent 20c1855
commit 703bead
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/cpp/pybind/docstring.cpp b/cpp/pybind/docstring.cpp
@@ -201,14 +201,24 @@ std::string FunctionDoc::ToGoogleDocString() const {
         for (size_t i = 0; i < overload.argument_docs_.size(); ++i) {
             const ArgumentDoc& argument_doc = overload.argument_docs_[i];
             rc << argument_doc.name_;
+            if (argument_doc.type_ != "") {
+                rc << ": " << argument_doc.type_;
+            }
             if (argument_doc.default_ != "") {
-                rc << "=" << argument_doc.default_;
+                rc << " = " << argument_doc.default_;
             }
             if (i != overload.argument_docs_.size() - 1) {
                 rc << ", ";
             }
         }
-        rc << ")" << std::endl;
+        rc << ")";
+
+        // Return type
+        if (overload.return_doc_.type_ != "") {
+            rc << " -> " << overload.return_doc_.type_;
+        }
+
+        rc << std::endl;
 
         // Summary line, strictly speaking this shall be at the very front.
         // However from a compiled Python module we need the function signature