Add support for static initialized defaults

Author: brianschubert

mypyc/codegen/emitclass.py

@@ -29,6 +29,7 @@
     FuncIR,
     get_text_signature,
 )
+from mypyc.ir.module_ir import ModuleIR, get_module_top_level
 from mypyc.ir.rtypes import RTuple, RType, object_rprimitive
 from mypyc.namegen import NameGenerator
 from mypyc.sametype import is_same_type
@@ -193,7 +194,7 @@ def generate_class_type_decl(
         )
 
 
-def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
+def generate_class(cl: ClassIR, module: ModuleIR, emitter: Emitter) -> None:
     """Generate C code for a class.
 
     This is the main entry point to the module.
@@ -333,7 +334,7 @@ def emit_line() -> None:
     if cl.is_trait:
         generate_new_for_trait(cl, new_name, emitter)
 
-    generate_methods_table(cl, methods_name, emitter)
+    generate_methods_table(cl, methods_name, emitter, module)
     emit_line()
 
     flags = ["Py_TPFLAGS_DEFAULT", "Py_TPFLAGS_HEAPTYPE", "Py_TPFLAGS_BASETYPE"]
@@ -352,7 +353,7 @@ def emit_line() -> None:
         flags.append("Py_TPFLAGS_MANAGED_DICT")
     fields["tp_flags"] = " | ".join(flags)
 
-    fields["tp_doc"] = native_class_doc_initializer(cl)
+    fields["tp_doc"] = native_class_doc_initializer(cl, get_module_top_level(module))
 
     emitter.emit_line(f"static PyTypeObject {emitter.type_struct_name(cl)}_template_ = {{")
     emitter.emit_line("PyVarObject_HEAD_INIT(NULL, 0)")
@@ -837,7 +838,7 @@ def generate_finalize_for_class(
     emitter.emit_line("}")
 
 
-def generate_methods_table(cl: ClassIR, name: str, emitter: Emitter) -> None:
+def generate_methods_table(cl: ClassIR, name: str, emitter: Emitter, module: ModuleIR) -> None:
     emitter.emit_line(f"static PyMethodDef {name}[] = {{")
     for fn in cl.methods.values():
         if fn.decl.is_prop_setter or fn.decl.is_prop_getter or fn.internal:
@@ -850,7 +851,7 @@ def generate_methods_table(cl: ClassIR, name: str, emitter: Emitter) -> None:
         elif fn.decl.kind == FUNC_CLASSMETHOD:
             flags.append("METH_CLASS")
 
-        doc = native_function_doc_initializer(fn)
+        doc = native_function_doc_initializer(fn, get_module_top_level(module))
         emitter.emit_line(" {}, {}}},".format(" | ".join(flags), doc))
 
     # Provide a default __getstate__ and __setstate__
@@ -1111,10 +1112,10 @@ def has_managed_dict(cl: ClassIR, emitter: Emitter) -> bool:
     )
 
 
-def native_class_doc_initializer(cl: ClassIR) -> str:
+def native_class_doc_initializer(cl: ClassIR, module_body: FuncIR) -> str:
     init_fn = cl.get_method("__init__")
     if init_fn is not None:
-        text_sig = get_text_signature(init_fn, bound=True)
+        text_sig = get_text_signature(init_fn, module_body, bound=True)
         if text_sig is None:
             return "NULL"
         text_sig = text_sig.replace("__init__", cl.name, 1)

mypyc/codegen/emitfunc.py

@@ -113,8 +113,8 @@ def native_function_header(fn: FuncDecl, emitter: Emitter) -> str:
     )
 
 
-def native_function_doc_initializer(func: FuncIR) -> str:
-    text_sig = get_text_signature(func)
+def native_function_doc_initializer(func: FuncIR, module_body: FuncIR) -> str:
+    text_sig = get_text_signature(func, module_body)
     if text_sig is None:
         return "NULL"
     docstring = f"{text_sig}\n--\n\n"

mypyc/codegen/emitmodule.py

@@ -54,7 +54,7 @@
 )
 from mypyc.errors import Errors
 from mypyc.ir.func_ir import FuncIR
-from mypyc.ir.module_ir import ModuleIR, ModuleIRs, deserialize_modules
+from mypyc.ir.module_ir import ModuleIR, ModuleIRs, deserialize_modules, get_module_top_level
 from mypyc.ir.ops import DeserMaps, LoadLiteral
 from mypyc.ir.rtypes import RType
 from mypyc.irbuild.main import build_ir
@@ -567,7 +567,7 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
 
             for cl in module.classes:
                 if cl.is_ext_class:
-                    generate_class(cl, module_name, emitter)
+                    generate_class(cl, module, emitter)
 
             # Generate Python extension module definitions and module initialization functions.
             self.generate_module_def(emitter, module_name, module)
@@ -919,7 +919,7 @@ def emit_module_methods(
                 flag = "METH_FASTCALL"
             else:
                 flag = "METH_VARARGS"
-            doc = native_function_doc_initializer(fn)
+            doc = native_function_doc_initializer(fn, get_module_top_level(module))
             emitter.emit_line(
                 (
                     '{{"{name}", (PyCFunction){prefix}{cname}, {flag} | METH_KEYWORDS, '

mypyc/ir/func_ir.py

@@ -13,13 +13,20 @@
     AssignMulti,
     BasicBlock,
     Box,
+    CallC,
     ControlOp,
     DeserMaps,
     Float,
+    GetElementPtr,
+    InitStatic,
     Integer,
+    IntOp,
     LoadAddress,
     LoadLiteral,
+    LoadMem,
+    LoadStatic,
     Register,
+    SetMem,
     TupleSet,
     Value,
 )
@@ -406,7 +413,7 @@ def all_values_full(args: list[Register], blocks: list[BasicBlock]) -> list[Valu
 _NOT_REPRESENTABLE = object()
 
 
-def get_text_signature(fn: FuncIR, *, bound: bool = False) -> str | None:
+def get_text_signature(fn: FuncIR, module_body: 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.
     """
@@ -430,7 +437,7 @@ def get_text_signature(fn: FuncIR, *, bound: bool = False) -> str | None:
         )
         default: object = inspect.Parameter.empty
         if arg.optional:
-            default = _find_default_argument(arg.name, fn.blocks)
+            default = _find_default_argument(arg.name, fn.blocks, module_body.blocks)
             if default is _NOT_REPRESENTABLE:
                 # This default argument cannot be represented in a __text_signature__
                 return None
@@ -444,16 +451,29 @@ def get_text_signature(fn: FuncIR, *, bound: bool = False) -> str | None:
     return f"{fn.name}{inspect.Signature(parameters)}"
 
 
-def _find_default_argument(name: str, blocks: list[BasicBlock]) -> object:
+def _find_default_argument(
+    name: str, blocks: list[BasicBlock], static_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)
+                if isinstance(op.src, LoadStatic):
+                    return _find_init_static(op.src.identifier, static_blocks)
+                else:
+                    return _extract_python_literal(op.src, blocks)
+    return _NOT_REPRESENTABLE
+
+
+def _find_init_static(fullname: str, blocks: list[BasicBlock]) -> object:
+    for block in blocks:
+        for op in block.ops:
+            if isinstance(op, InitStatic) and op.identifier == fullname:
+                return _extract_python_literal(op.value, blocks)
     return _NOT_REPRESENTABLE
 
 
-def _extract_python_literal(value: Value) -> object:
+def _extract_python_literal(value: Value, blocks: list[BasicBlock]) -> object:
     if isinstance(value, Integer):
         if is_none_rprimitive(value.type):
             return None
@@ -466,10 +486,66 @@ def _extract_python_literal(value: Value) -> object:
     elif isinstance(value, LoadLiteral):
         return value.value
     elif isinstance(value, Box):
-        return _extract_python_literal(value.src)
+        return _extract_python_literal(value.src, blocks)
     elif isinstance(value, TupleSet):
-        items = tuple(_extract_python_literal(item) for item in value.items)
+        items = tuple(_extract_python_literal(item, blocks) for item in value.items)
         if any(itm is _NOT_REPRESENTABLE for itm in items):
             return _NOT_REPRESENTABLE
         return items
+    elif isinstance(value, CallC):
+        if value.function_name == "PyList_New":
+            assert len(value.args) == 1
+            size = _extract_python_literal(value.args[0], blocks)
+            if size == 0:
+                return []
+            return _extract_list(value, blocks)
+        if value.function_name == "PyDict_New":
+            return {}
+        if value.function_name == "CPyDict_Build":
+            args = [_extract_python_literal(arg, blocks) for arg in value.args]
+            if any(arg is _NOT_REPRESENTABLE for arg in args):
+                return _NOT_REPRESENTABLE
+            return {k: v for k, v in zip(args[1::2], args[2::2])}
+        if value.function_name == "PySet_New":
+            result = _extract_set(value, blocks)
+            if not result:
+                return _NOT_REPRESENTABLE  # set() isn't valid in __text_signature__
+            return result
+    elif isinstance(value, LoadAddress) and value.src == "_Py_EllipsisObject":
+        return _EllipsisLiteral()
     return _NOT_REPRESENTABLE
+
+
+def _extract_list(value: CallC, blocks: list[BasicBlock]) -> object:
+    result = []
+    for block in blocks:
+        for op in block.ops:
+            if isinstance(op, SetMem):
+                dest = op.dest.lhs if isinstance(op.dest, IntOp) else op.dest
+                if (
+                    isinstance(dest, LoadMem)
+                    and isinstance(dest.src, GetElementPtr)
+                    and dest.src.src == value
+                ):
+                    item = _extract_python_literal(op.src, blocks)
+                    if item is _NOT_REPRESENTABLE:
+                        return _NOT_REPRESENTABLE
+                    result.append(item)
+    return result
+
+
+def _extract_set(value: CallC, blocks: list[BasicBlock]) -> object:
+    result = set()
+    for block in blocks:
+        for op in block.ops:
+            if isinstance(op, CallC) and op.function_name == "PySet_Add" and op.args[0] == value:
+                item = _extract_python_literal(op.args[1], blocks)
+                if item is _NOT_REPRESENTABLE:
+                    return _NOT_REPRESENTABLE
+                result.add(item)
+    return result
+
+
+class _EllipsisLiteral:
+    def __repr__(self) -> str:
+        return "..."

mypyc/ir/module_ir.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from mypyc.common import JsonDict
+from mypyc.common import TOP_LEVEL_NAME, JsonDict
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR
 from mypyc.ir.ops import DeserMaps
@@ -90,3 +90,11 @@ def deserialize_modules(data: dict[str, JsonDict], ctx: DeserMaps) -> dict[str,
 # ModulesIRs should also always be an *OrderedDict*, but if we
 # declared it that way we would need to put it in quotes everywhere...
 ModuleIRs = dict[str, ModuleIR]
+
+
+def get_module_top_level(module: ModuleIR) -> FuncIR:
+    # Optimization: we tend to put the top level last, so reverse iterate
+    for fn in reversed(module.functions):
+        if fn.name == TOP_LEVEL_NAME:
+            return fn
+    assert False, f"module '{module.fullname}' missing '{TOP_LEVEL_NAME}' function"

mypyc/test-data/run-signatures.test

@@ -42,6 +42,16 @@ def default_tuple_empty(x=()): pass
 def default_tuple_literals(x=(1, "a", 1.0, False, True, None, (), (1,2,(3,4)))): pass
 def default_tuple_singleton(x=(1,)): pass
 def default_named_constant(x=A): pass
+def default_complex(x=1+2j): pass
+def default_list_empty(x=[]): pass
+def default_list_literals(x=[1, 2, 3]): pass
+def default_dict_empty(x={}): pass
+def default_dict_literals(x={'a': 1}): pass
+def default_set_literals(x={1, 2, 3}): pass
+def default_tuple_collections(x=([1], {2}, {3: 4})): pass
+def default_nested_collections(x={'list': [1], 'set': {2}, 'dict': {3: 4}}): pass
+def default_literal_fold(x=1+2): pass
+def default_ellipsis(x=...): pass
 
 [file driver.py]
 import inspect
@@ -56,6 +66,16 @@ assert str(inspect.signature(default_none)) == "(x=None)"
 assert str(inspect.signature(default_tuple_empty)) == "(x=())"
 assert str(inspect.signature(default_tuple_literals)) == "(x=(1, 'a', 1.0, False, True, None, (), (1, 2, (3, 4))))"
 assert str(inspect.signature(default_named_constant)) == "(x=1)"
+assert str(inspect.signature(default_complex)) == "(x=(1+2j))"
+assert str(inspect.signature(default_list_empty)) == "(x=[])"
+assert str(inspect.signature(default_list_literals)) == "(x=[1, 2, 3])"
+assert str(inspect.signature(default_dict_empty)) == "(x={})"
+assert str(inspect.signature(default_dict_literals)) == "(x={'a': 1})"
+assert str(inspect.signature(default_set_literals)) == "(x={1, 2, 3})"
+assert str(inspect.signature(default_tuple_collections)) == "(x=([1], {2}, {3: 4}))"
+assert str(inspect.signature(default_nested_collections)) == "(x={'list': [1], 'set': {2}, 'dict': {3: 4}})"
+assert str(inspect.signature(default_literal_fold)) == "(x=3)"
+assert str(inspect.signature(default_ellipsis)) == "(x=Ellipsis)"
 
 # Check __text_signature__ directly since inspect.signature produces
 # an incorrect signature for 1-tuple default arguments prior to
@@ -95,19 +115,6 @@ def bad_set_empty(x=set()): pass  # supported by ast.literal_eval, but not by in
 def bad_nan(x=float("nan")): pass
 def bad_enum(x=Color.RED): pass
 
-# TODO: Default arguments that could potentially be represented in a
-# __text_signature__, but which are not currently supported.
-# See 'inspect._signature_fromstr' for what default values are supported at runtime.
-def bad_complex(x=1+2j): pass
-def bad_list_empty(x=[]): pass
-def bad_list_literals(x=[1, 2, 3]): pass
-def bad_dict_empty(x={}): pass
-def bad_dict_literals(x={'a': 1}): pass
-def bad_set_literals(x={1, 2, 3}): pass
-def bad_tuple_literals(x=([1, 2, 3], {'a': 1}, {1, 2, 3})): pass
-def bad_ellipsis(x=...): pass
-def bad_literal_fold(x=1+2): pass
-
 [file driver.py]
 import inspect
 from testutil import assertRaises
@@ -155,7 +162,7 @@ class HasInit:
 class InheritedInit(HasInit): pass
 
 class HasInitBad:
-    def __init__(self, x=[]) -> None: pass
+    def __init__(self, x=object()) -> None: pass
 
 [file driver.py]
 import inspect