[mypyc] Generate introspection signatures for compiled functions (#19307)

Refs mypyc/mypyc#838

This PR populates `__text_signature__` for compiled functions, making
runtime signature introspection possible (i.e.
`inspect.signature(func)`).

While `__text_signature__` is an undocumented CPython implementation
detail, other extension module generators are using it in practice. For
example, PyO3 and Cython both support it. I think it would be reasonable
for mypyc to support it too.

Some function signatures can't be represented by `__text_signature__`
(namely, those with complex default arguments). In those cases, no
signatures are generated (same as the current behavior).

Author: brianschubert

mypyc/codegen/emitclass.py

@@ -5,8 +5,9 @@
 from collections.abc import Mapping
 from typing import Callable
 
+from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import Emitter, HeaderDeclaration, ReturnHandler
-from mypyc.codegen.emitfunc import native_function_header
+from mypyc.codegen.emitfunc import native_function_doc_initializer, native_function_header
 from mypyc.codegen.emitwrapper import (
     generate_bin_op_wrapper,
     generate_bool_wrapper,
@@ -21,7 +22,13 @@
 )
 from mypyc.common import BITMAP_BITS, BITMAP_TYPE, NATIVE_PREFIX, PREFIX, REG_PREFIX
 from mypyc.ir.class_ir import ClassIR, VTableEntries
-from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR
+from mypyc.ir.func_ir import (
+    FUNC_CLASSMETHOD,
+    FUNC_STATICMETHOD,
+    FuncDecl,
+    FuncIR,
+    get_text_signature,
+)
 from mypyc.ir.rtypes import RTuple, RType, object_rprimitive
 from mypyc.namegen import NameGenerator
 from mypyc.sametype import is_same_type
@@ -368,6 +375,8 @@ def emit_line() -> None:
         flags.append("Py_TPFLAGS_MANAGED_DICT")
     fields["tp_flags"] = " | ".join(flags)
 
+    fields["tp_doc"] = native_class_doc_initializer(cl)
+
     emitter.emit_line(f"static PyTypeObject {emitter.type_struct_name(cl)}_template_ = {{")
     emitter.emit_line("PyVarObject_HEAD_INIT(NULL, 0)")
     for field, value in fields.items():
@@ -915,7 +924,8 @@ def generate_methods_table(cl: ClassIR, name: str, emitter: Emitter) -> None:
         elif fn.decl.kind == FUNC_CLASSMETHOD:
             flags.append("METH_CLASS")
 
-        emitter.emit_line(" {}, NULL}},".format(" | ".join(flags)))
+        doc = native_function_doc_initializer(fn)
+        emitter.emit_line(" {}, {}}},".format(" | ".join(flags), doc))
 
     # Provide a default __getstate__ and __setstate__
     if not cl.has_method("__setstate__") and not cl.has_method("__getstate__"):
@@ -1173,3 +1183,16 @@ def has_managed_dict(cl: ClassIR, emitter: Emitter) -> bool:
         and cl.has_dict
         and cl.builtin_base != "PyBaseExceptionObject"
     )
+
+
+def native_class_doc_initializer(cl: ClassIR) -> str:
+    init_fn = cl.get_method("__init__")
+    if init_fn is not None:
+        text_sig = get_text_signature(init_fn, bound=True)
+        if text_sig is None:
+            return "NULL"
+        text_sig = text_sig.replace("__init__", cl.name, 1)
+    else:
+        text_sig = f"{cl.name}()"
+    docstring = f"{text_sig}\n--\n\n"
+    return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))

mypyc/codegen/emitfunc.py

@@ -5,6 +5,7 @@
 from typing import Final
 
 from mypyc.analysis.blockfreq import frequently_executed_blocks
+from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import DEBUG_ERRORS, Emitter, TracebackAndGotoHandler, c_array_initializer
 from mypyc.common import (
     HAVE_IMMORTAL,
@@ -16,7 +17,14 @@
     TYPE_VAR_PREFIX,
 )
 from mypyc.ir.class_ir import ClassIR
-from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD, FuncDecl, FuncIR, all_values
+from mypyc.ir.func_ir import (
+    FUNC_CLASSMETHOD,
+    FUNC_STATICMETHOD,
+    FuncDecl,
+    FuncIR,
+    all_values,
+    get_text_signature,
+)
 from mypyc.ir.ops import (
     ERR_FALSE,
     NAMESPACE_MODULE,
@@ -106,6 +114,14 @@ def native_function_header(fn: FuncDecl, emitter: Emitter) -> str:
     )
 
 
+def native_function_doc_initializer(func: FuncIR) -> str:
+    text_sig = get_text_signature(func)
+    if text_sig is None:
+        return "NULL"
+    docstring = f"{text_sig}\n--\n\n"
+    return c_string_initializer(docstring.encode("ascii", errors="backslashreplace"))
+
+
 def generate_native_function(
     fn: FuncIR, emitter: Emitter, source_path: str, module_name: str
 ) -> None:

mypyc/codegen/emitmodule.py

@@ -30,7 +30,11 @@
 from mypyc.codegen.cstring import c_string_initializer
 from mypyc.codegen.emit import Emitter, EmitterContext, HeaderDeclaration, c_array_initializer
 from mypyc.codegen.emitclass import generate_class, generate_class_reuse, generate_class_type_decl
-from mypyc.codegen.emitfunc import generate_native_function, native_function_header
+from mypyc.codegen.emitfunc import (
+    generate_native_function,
+    native_function_doc_initializer,
+    native_function_header,
+)
 from mypyc.codegen.emitwrapper import (
     generate_legacy_wrapper_function,
     generate_wrapper_function,
@@ -917,11 +921,14 @@ def emit_module_methods(
                 flag = "METH_FASTCALL"
             else:
                 flag = "METH_VARARGS"
+            doc = native_function_doc_initializer(fn)
             emitter.emit_line(
                 (
                     '{{"{name}", (PyCFunction){prefix}{cname}, {flag} | METH_KEYWORDS, '
-                    "NULL /* docstring */}},"
-                ).format(name=name, cname=fn.cname(emitter.names), prefix=PREFIX, flag=flag)
+                    "{doc} /* docstring */}},"
+                ).format(
+                    name=name, cname=fn.cname(emitter.names), prefix=PREFIX, flag=flag, doc=doc
+                )
             )
         emitter.emit_line("{NULL, NULL, 0, NULL}")
         emitter.emit_line("};")

mypyc/doc/differences_from_python.rst

@@ -316,7 +316,8 @@ non-exhaustive list of what won't work:
 - Instance ``__annotations__`` is usually not kept
 - Frames of compiled functions can't be inspected using ``inspect``
 - Compiled methods aren't considered methods by ``inspect.ismethod``
-- ``inspect.signature`` chokes on compiled functions
+- ``inspect.signature`` chokes on compiled functions with default arguments that
+  are not simple literals
 - ``inspect.iscoroutinefunction`` and ``asyncio.iscoroutinefunction`` will always return False for compiled functions, even those defined with `async def`
 
 Profiling hooks and tracing

mypyc/ir/func_ir.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 from collections.abc import Sequence
 from typing import Final
 
@@ -11,13 +12,24 @@
     Assign,
     AssignMulti,
     BasicBlock,
+    Box,
     ControlOp,
     DeserMaps,
+    Float,
+    Integer,
     LoadAddress,
+    LoadLiteral,
     Register,
+    TupleSet,
     Value,
 )
-from mypyc.ir.rtypes import RType, bitmap_rprimitive, deserialize_type
+from mypyc.ir.rtypes import (
+    RType,
+    bitmap_rprimitive,
+    deserialize_type,
+    is_bool_rprimitive,
+    is_none_rprimitive,
+)
 from mypyc.namegen import NameGenerator
 
 
@@ -379,3 +391,85 @@ def all_values_full(args: list[Register], blocks: list[BasicBlock]) -> list[Valu
                     values.append(op)
 
     return values
+
+
+_ARG_KIND_TO_INSPECT: Final = {
+    ArgKind.ARG_POS: inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    ArgKind.ARG_OPT: inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    ArgKind.ARG_STAR: inspect.Parameter.VAR_POSITIONAL,
+    ArgKind.ARG_NAMED: inspect.Parameter.KEYWORD_ONLY,
+    ArgKind.ARG_STAR2: inspect.Parameter.VAR_KEYWORD,
+    ArgKind.ARG_NAMED_OPT: inspect.Parameter.KEYWORD_ONLY,
+}
+
+# Sentinel indicating a value that cannot be represented in a text signature.
+_NOT_REPRESENTABLE = object()
+
+
+def get_text_signature(fn: FuncIR, *, bound: bool = False) -> str | None:
+    """Return a text signature in CPython's internal doc format, or None
+    if the function's signature cannot be represented.
+    """
+    parameters = []
+    mark_self = (fn.class_name is not None) and (fn.decl.kind != FUNC_STATICMETHOD) and not bound
+    sig = fn.decl.bound_sig if bound and fn.decl.bound_sig is not None else fn.decl.sig
+    # Pre-scan for end of positional-only parameters.
+    # This is needed to handle signatures like 'def foo(self, __x)', where mypy
+    # currently sees 'self' as being positional-or-keyword and '__x' as positional-only.
+    pos_only_idx = -1
+    for idx, arg in enumerate(sig.args):
+        if arg.pos_only and arg.kind in (ArgKind.ARG_POS, ArgKind.ARG_OPT):
+            pos_only_idx = idx
+    for idx, arg in enumerate(sig.args):
+        if arg.name.startswith(("__bitmap", "__mypyc")):
+            continue
+        kind = (
+            inspect.Parameter.POSITIONAL_ONLY
+            if idx <= pos_only_idx
+            else _ARG_KIND_TO_INSPECT[arg.kind]
+        )
+        default: object = inspect.Parameter.empty
+        if arg.optional:
+            default = _find_default_argument(arg.name, fn.blocks)
+            if default is _NOT_REPRESENTABLE:
+                # This default argument cannot be represented in a __text_signature__
+                return None
+
+        curr_param = inspect.Parameter(arg.name, kind, default=default)
+        parameters.append(curr_param)
+        if mark_self:
+            # Parameter.__init__/Parameter.replace do not accept $
+            curr_param._name = f"${arg.name}"  # type: ignore[attr-defined]
+            mark_self = False
+    return f"{fn.name}{inspect.Signature(parameters)}"
+
+
+def _find_default_argument(name: str, blocks: list[BasicBlock]) -> object:
+    # Find assignment inserted by gen_arg_defaults. Assumed to be the first assignment.
+    for block in blocks:
+        for op in block.ops:
+            if isinstance(op, Assign) and op.dest.name == name:
+                return _extract_python_literal(op.src)
+    return _NOT_REPRESENTABLE
+
+
+def _extract_python_literal(value: Value) -> object:
+    if isinstance(value, Integer):
+        if is_none_rprimitive(value.type):
+            return None
+        val = value.numeric_value()
+        if is_bool_rprimitive(value.type):
+            return bool(val)
+        return val
+    elif isinstance(value, Float):
+        return value.value
+    elif isinstance(value, LoadLiteral):
+        return value.value
+    elif isinstance(value, Box):
+        return _extract_python_literal(value.src)
+    elif isinstance(value, TupleSet):
+        items = tuple(_extract_python_literal(item) for item in value.items)
+        if any(itm is _NOT_REPRESENTABLE for itm in items):
+            return _NOT_REPRESENTABLE
+        return items
+    return _NOT_REPRESENTABLE