Guess module name for boost python types

katana/stubgen_katana.py

@@ -26,15 +26,6 @@
 DISABLED_CODES = '# mypy: disable-error-code="misc, override, attr-defined, no-redef, assignment"\n\n'
 
 
-def is_valid(type_name: str) -> bool:
-    try:
-        parse_type_comment(type_name, 0, 0, None)
-    except Exception:
-        return False
-    else:
-        return True
-
-
 def cleanup_type(type_name: str) -> str:
     type_name = type_name.replace('\n', ' ')
     type_name = type_name.rstrip('.')
@@ -124,26 +115,8 @@ def numeric_tuple_sub(m):
 
 
 class KatanaDocstringSignatureGenerator(DocstringSignatureGenerator):
-    def cleanup_sig(self, sig: FunctionSig) -> tuple[FunctionSig, list[str]]:
-        args = []
-        return_type = None
-        invalid = []
-        for arg in sig.args:
-            type_name = None
-            if arg.type:
-                type_name = cleanup_type(arg.type)
-                if not is_valid(type_name):
-                    invalid.append("  Invalid arg {}: {} {}".format(
-                        arg.name, repr(arg.type), repr(type_name)))
-                    type_name = None
-            args.append(ArgSig(arg.name, type_name, arg.default))
-        if sig.ret_type:
-            return_type = cleanup_type(sig.ret_type)
-            if not is_valid(return_type):
-                invalid.append("  Invalid ret: {} {}".format(
-                    repr(sig.ret_type), repr(return_type)))
-                return_type = None
-        return FunctionSig(sig.name, args, return_type), invalid
+    def cleanup_type(self, type_name: str) -> str:
+        return cleanup_type(type_name)
 
 
 class KatanaCSignatureGenerator(CDocstringSignatureGenerator):

stubgenlib.py

@@ -1,14 +1,59 @@
 from __future__ import absolute_import, annotations, division, print_function
 
+import re
 from typing import Any
 
-import docstring_parser
+from mypy.stubdoc import infer_sig_from_docstring
 from mypy.stubgenc import (
     ArgSig, FunctionContext, FunctionSig, SignatureGenerator
 )
+from mypy.fastparse import parse_type_comment
 
 
-class DocstringSignatureGenerator(SignatureGenerator):
+class BaseSigFixer:
+
+    @staticmethod
+    def is_valid(type_name: str) -> bool:
+        try:
+            parse_type_comment(type_name, 0, 0, None)
+        except Exception:
+            return False
+        else:
+            return True
+
+    def cleanup_type(self, type_name: str, ctx: FunctionContext, is_result: bool) -> str:
+        """Override this to implement logic to fix a type"""
+        return type_name
+
+    def cleanup_sig_types(self, sig: FunctionSig, ctx: FunctionContext) -> FunctionSig:
+        args = []
+        return_type = None
+        invalid = []
+        for arg in sig.args:
+            type_name = None
+            if arg.type:
+                type_name = self.cleanup_type(arg.type, ctx, is_result=False)
+                if not self.is_valid(type_name):
+                    invalid.append("  Invalid arg {}: {} {}".format(
+                        arg.name, repr(arg.type), repr(type_name)))
+                    type_name = None
+            args.append(ArgSig(arg.name, type_name, arg.default))
+        if sig.ret_type:
+            return_type = self.cleanup_type(sig.ret_type, ctx, is_result=True)
+            if not self.is_valid(return_type):
+                invalid.append("  Invalid ret: {} {}".format(
+                    repr(sig.ret_type), repr(return_type)))
+                return_type = None
+
+        if invalid:
+            print(f"Invalid type after cleanup: {ctx.fullname}")
+            print(sig.format_sig())
+
+        # FIXME: only copy if something has changed?
+        return FunctionSig(sig.name, args, return_type)
+
+
+class DocstringSignatureGenerator(SignatureGenerator, BaseSigFixer):
     """
     Generate signatures from docstrings.
 
@@ -17,12 +62,10 @@ class DocstringSignatureGenerator(SignatureGenerator):
     google, and epydoc (which Katana uses).
     """
 
-    def cleanup_sig(self, sig):
-        return sig, []
-
     def get_function_sig(
         self, default_sig: FunctionSig, ctx: FunctionContext
     ) -> list[FunctionSig] | None:
+        import docstring_parser
         if ctx.docstr:
             parsed = docstring_parser.parse(ctx.docstr)
             args = []
@@ -37,14 +80,26 @@ def get_function_sig(
             if parsed.returns and parsed.returns.type_name:
                 return_type = parsed.returns.type_name
             sig = FunctionSig(ctx.name, args, return_type)
-            sig, invalid = self.cleanup_sig(sig)
-            if invalid:
-                print(ctx.fullname)
-                print(sig.format_sig())
+            sig = self.cleanup_sig_types(sig, ctx)
+
             # if len(default_sig.args) != len(args):
             #     print("  Lengths are not the same")
 
             merged_sig = default_sig.merge(sig)
             # print(merged_sig.format_sig())
             return [merged_sig]
         return None
+
+
+class BoostDocstringSignatureGenerator(SignatureGenerator):
+    def get_function_sig(
+        self, default_sig: FunctionSig, ctx: FunctionContext
+    ) -> list[FunctionSig] | None:
+
+        if ctx.docstr:
+            # convert the boost-provided signature into a proper python signature
+            docstr = ctx.docstr.replace('[', '')
+            docstr = docstr.replace(']', '')
+            docstr = re.sub(r"\(([^(]+)\)([a-zA-Z_][a-zA-Z0-9_]*)", lambda m: '{}: {}'.format(m.group(2), m.group(1)), docstr)
+            return infer_sig_from_docstring(docstr, ctx.name)
+