[mypyc] Support __setitem__, __delitem__, __len__ and __contains__ (#10451)

Previously these couldn't be used in compiled classes.

__setitem__ and __delitem__ are implemented by a single wrapper function. If the value is
NULL, it should act as __delitem__. This makes the implementation non-trivial. First, we
need to support both a single variant and two variants being defined in a class. Second, when
overriding, we need to use super(), if we only override one of them.

The other dunders are straightforward.

The implementation is pretty verbose. I'll look at refactoring it in a separate PR.

This has some overlap with the previous PR #10211 by @sohailsomani.

Author: JukkaL

mypyc/codegen/emitclass.py

@@ -9,7 +9,8 @@
 from mypyc.codegen.emitfunc import native_function_header
 from mypyc.codegen.emitwrapper import (
     generate_dunder_wrapper, generate_hash_wrapper, generate_richcompare_wrapper,
-    generate_bool_wrapper, generate_get_wrapper,
+    generate_bool_wrapper, generate_get_wrapper, generate_len_wrapper,
+    generate_set_del_item_wrapper, generate_contains_wrapper
 )
 from mypyc.ir.rtypes import RType, RTuple, object_rprimitive
 from mypyc.ir.func_ir import FuncIR, FuncDecl, FUNC_STATICMETHOD, FUNC_CLASSMETHOD
@@ -46,6 +47,13 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 
 AS_MAPPING_SLOT_DEFS = {
     '__getitem__': ('mp_subscript', generate_dunder_wrapper),
+    '__setitem__': ('mp_ass_subscript', generate_set_del_item_wrapper),
+    '__delitem__': ('mp_ass_subscript', generate_set_del_item_wrapper),
+    '__len__': ('mp_length', generate_len_wrapper),
+}  # type: SlotTable
+
+AS_SEQUENCE_SLOT_DEFS = {
+    '__contains__': ('sq_contains', generate_contains_wrapper),
 }  # type: SlotTable
 
 AS_NUMBER_SLOT_DEFS = {
@@ -60,6 +68,7 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 
 SIDE_TABLES = [
     ('as_mapping', 'PyMappingMethods', AS_MAPPING_SLOT_DEFS),
+    ('as_sequence', 'PySequenceMethods', AS_SEQUENCE_SLOT_DEFS),
     ('as_number', 'PyNumberMethods', AS_NUMBER_SLOT_DEFS),
     ('as_async', 'PyAsyncMethods', AS_ASYNC_SLOT_DEFS),
 ]
@@ -82,11 +91,19 @@ def generate_call_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 
 def generate_slots(cl: ClassIR, table: SlotTable, emitter: Emitter) -> Dict[str, str]:
     fields = OrderedDict()  # type: Dict[str, str]
+    generated = {}  # type: Dict[str, str]
     # Sort for determinism on Python 3.5
-    for name, (slot, generator) in sorted(table.items()):
+    for name, (slot, generator) in sorted(table.items(), reverse=True):
         method_cls = cl.get_method_and_class(name)
         if method_cls and (method_cls[1] == cl or name in ALWAYS_FILL):
-            fields[slot] = generator(cl, method_cls[0], emitter)
+            if slot in generated:
+                # Reuse previously generated wrapper.
+                fields[slot] = generated[slot]
+            else:
+                # Generate new wrapper.
+                name = generator(cl, method_cls[0], emitter)
+                fields[slot] = name
+                generated[slot] = name
 
     return fields
 

mypyc/codegen/emitwrapper.py

@@ -10,7 +10,7 @@
 or methods in a single compilation unit.
 """
 
-from typing import List, Optional
+from typing import List, Optional, Sequence
 
 from mypy.nodes import ARG_POS, ARG_OPT, ARG_NAMED_OPT, ARG_NAMED, ARG_STAR, ARG_STAR2
 
@@ -350,6 +350,29 @@ def generate_hash_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     return name
 
 
+def generate_len_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    """Generates a wrapper for native __len__ methods."""
+    name = '{}{}{}'.format(DUNDER_PREFIX, fn.name, cl.name_prefix(emitter.names))
+    emitter.emit_line('static Py_ssize_t {name}(PyObject *self) {{'.format(
+        name=name
+    ))
+    emitter.emit_line('{}retval = {}{}{}(self);'.format(emitter.ctype_spaced(fn.ret_type),
+                                                        emitter.get_group_prefix(fn.decl),
+                                                        NATIVE_PREFIX,
+                                                        fn.cname(emitter.names)))
+    emitter.emit_error_check('retval', fn.ret_type, 'return -1;')
+    if is_int_rprimitive(fn.ret_type):
+        emitter.emit_line('Py_ssize_t val = CPyTagged_AsSsize_t(retval);')
+    else:
+        emitter.emit_line('Py_ssize_t val = PyLong_AsSsize_t(retval);')
+    emitter.emit_dec_ref('retval', fn.ret_type)
+    emitter.emit_line('if (PyErr_Occurred()) return -1;')
+    emitter.emit_line('return val;')
+    emitter.emit_line('}')
+
+    return name
+
+
 def generate_bool_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     """Generates a wrapper for native __bool__ methods."""
     name = '{}{}{}'.format(DUNDER_PREFIX, fn.name, cl.name_prefix(emitter.names))
@@ -370,6 +393,134 @@ def generate_bool_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     return name
 
 
+def generate_del_item_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    """Generates a wrapper for native __delitem__.
+
+    This is only called from a combined __delitem__/__setitem__ wrapper.
+    """
+    name = '{}{}{}'.format(DUNDER_PREFIX, '__delitem__', cl.name_prefix(emitter.names))
+    input_args = ', '.join('PyObject *obj_{}'.format(arg.name) for arg in fn.args)
+    emitter.emit_line('static int {name}({input_args}) {{'.format(
+        name=name,
+        input_args=input_args,
+    ))
+    generate_set_del_item_wrapper_inner(fn, emitter, fn.args)
+    return name
+
+
+def generate_set_del_item_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    """Generates a wrapper for native __setitem__ method (also works for __delitem__).
+
+    This is used with the mapping protocol slot. Arguments are taken as *PyObjects and we
+    return a negative C int on error.
+
+    Create a separate wrapper function for __delitem__ as needed and have the
+    __setitem__ wrapper call it if the value is NULL. Return the name
+    of the outer (__setitem__) wrapper.
+    """
+    method_cls = cl.get_method_and_class('__delitem__')
+    del_name = None
+    if method_cls and method_cls[1] == cl:
+        # Generate a separate wrapper for __delitem__
+        del_name = generate_del_item_wrapper(cl, method_cls[0], emitter)
+
+    args = fn.args
+    if fn.name == '__delitem__':
+        # Add an extra argument for value that we expect to be NULL.
+        args = list(args) + [RuntimeArg('___value', object_rprimitive, ARG_POS)]
+
+    name = '{}{}{}'.format(DUNDER_PREFIX, '__setitem__', cl.name_prefix(emitter.names))
+    input_args = ', '.join('PyObject *obj_{}'.format(arg.name) for arg in args)
+    emitter.emit_line('static int {name}({input_args}) {{'.format(
+        name=name,
+        input_args=input_args,
+    ))
+
+    # First check if this is __delitem__
+    emitter.emit_line('if (obj_{} == NULL) {{'.format(args[2].name))
+    if del_name is not None:
+        # We have a native implementation, so call it
+        emitter.emit_line('return {}(obj_{}, obj_{});'.format(del_name,
+                                                              args[0].name,
+                                                              args[1].name))
+    else:
+        # Try to call superclass method instead
+        emitter.emit_line(
+            'PyObject *super = CPy_Super(CPyModule_builtins, obj_{});'.format(args[0].name))
+        emitter.emit_line('if (super == NULL) return -1;')
+        emitter.emit_line(
+            'PyObject *result = PyObject_CallMethod(super, "__delitem__", "O", obj_{});'.format(
+                args[1].name))
+        emitter.emit_line('Py_DECREF(super);')
+        emitter.emit_line('Py_XDECREF(result);')
+        emitter.emit_line('return result == NULL ? -1 : 0;')
+    emitter.emit_line('}')
+
+    method_cls = cl.get_method_and_class('__setitem__')
+    if method_cls and method_cls[1] == cl:
+        generate_set_del_item_wrapper_inner(fn, emitter, args)
+    else:
+        emitter.emit_line(
+            'PyObject *super = CPy_Super(CPyModule_builtins, obj_{});'.format(args[0].name))
+        emitter.emit_line('if (super == NULL) return -1;')
+        emitter.emit_line('PyObject *result;')
+
+        if method_cls is None and cl.builtin_base is None:
+            msg = "'{}' object does not support item assignment".format(cl.name)
+            emitter.emit_line(
+                'PyErr_SetString(PyExc_TypeError, "{}");'.format(msg))
+            emitter.emit_line('result = NULL;')
+        else:
+            # A base class may have __setitem__
+            emitter.emit_line(
+                'result = PyObject_CallMethod(super, "__setitem__", "OO", obj_{}, obj_{});'.format(
+                    args[1].name, args[2].name))
+        emitter.emit_line('Py_DECREF(super);')
+        emitter.emit_line('Py_XDECREF(result);')
+        emitter.emit_line('return result == NULL ? -1 : 0;')
+        emitter.emit_line('}')
+    return name
+
+
+def generate_set_del_item_wrapper_inner(fn: FuncIR, emitter: Emitter,
+                                        args: Sequence[RuntimeArg]) -> None:
+    for arg in args:
+        generate_arg_check(arg.name, arg.type, emitter, 'goto fail;', False)
+    native_args = ', '.join('arg_{}'.format(arg.name) for arg in args)
+    emitter.emit_line('{}val = {}{}({});'.format(emitter.ctype_spaced(fn.ret_type),
+                                                 NATIVE_PREFIX,
+                                                 fn.cname(emitter.names),
+                                                 native_args))
+    emitter.emit_error_check('val', fn.ret_type, 'goto fail;')
+    emitter.emit_dec_ref('val', fn.ret_type)
+    emitter.emit_line('return 0;')
+    emitter.emit_label('fail')
+    emitter.emit_line('return -1;')
+    emitter.emit_line('}')
+
+
+def generate_contains_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
+    """Generates a wrapper for a native __contains__ method."""
+    name = '{}{}{}'.format(DUNDER_PREFIX, fn.name, cl.name_prefix(emitter.names))
+    emitter.emit_line(
+        'static int {name}(PyObject *self, PyObject *obj_item) {{'.
+        format(name=name))
+    generate_arg_check('item', fn.args[1].type, emitter, 'return -1;', False)
+    emitter.emit_line('{}val = {}{}(self, arg_item);'.format(emitter.ctype_spaced(fn.ret_type),
+                                                             NATIVE_PREFIX,
+                                                             fn.cname(emitter.names)))
+    emitter.emit_error_check('val', fn.ret_type, 'return -1;')
+    if is_bool_rprimitive(fn.ret_type):
+        emitter.emit_line('return val;')
+    else:
+        emitter.emit_line('int boolval = PyObject_IsTrue(val);')
+        emitter.emit_dec_ref('val', fn.ret_type)
+        emitter.emit_line('return boolval;')
+    emitter.emit_line('}')
+
+    return name
+
+
 # Helpers
 
 

mypyc/irbuild/ll_builder.py

@@ -646,6 +646,8 @@ def binary_op(self,
         if is_bool_rprimitive(ltype) and is_bool_rprimitive(rtype) and op in (
                 '&', '&=', '|', '|=', '^', '^='):
             return self.bool_bitwise_op(lreg, rreg, op[0], line)
+        if isinstance(rtype, RInstance) and op in ('in', 'not in'):
+            return self.translate_instance_contains(rreg, lreg, op, line)
 
         call_c_ops_candidates = binary_ops.get(op, [])
         target = self.matching_call_c(call_c_ops_candidates, [lreg, rreg], line)
@@ -829,6 +831,14 @@ def compare_tuples(self,
         self.goto_and_activate(out)
         return result
 
+    def translate_instance_contains(self, inst: Value, item: Value, op: str, line: int) -> Value:
+        res = self.gen_method_call(inst, '__contains__', [item], None, line)
+        if not is_bool_rprimitive(res.type):
+            res = self.call_c(bool_op, [res], line)
+        if op == 'not in':
+            res = self.bool_bitwise_op(res, Integer(1, rtype=bool_rprimitive), '^', line)
+        return res
+
     def bool_bitwise_op(self, lreg: Value, rreg: Value, op: str, line: int) -> Value:
         if op == '&':
             code = IntOp.AND
@@ -1099,9 +1109,12 @@ def int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int) -> Val
     def comparison_op(self, lhs: Value, rhs: Value, op: int, line: int) -> Value:
         return self.add(ComparisonOp(lhs, rhs, op, line))
 
-    def builtin_len(self, val: Value, line: int,
-                    use_pyssize_t: bool = False) -> Value:
-        """Return short_int_rprimitive by default."""
+    def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Value:
+        """Generate len(val).
+
+        Return short_int_rprimitive by default.
+        Return c_pyssize_t if use_pyssize_t is true (unshifted).
+        """
         typ = val.type
         if is_list_rprimitive(typ) or is_tuple_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PyVarObject, 'ob_size'))
@@ -1128,8 +1141,22 @@ def builtin_len(self, val: Value, line: int,
             offset = Integer(1, c_pyssize_t_rprimitive, line)
             return self.int_op(short_int_rprimitive, size_value, offset,
                                IntOp.LEFT_SHIFT, line)
-        # generic case
+        elif isinstance(typ, RInstance):
+            # TODO: Support use_pyssize_t
+            assert not use_pyssize_t
+            length = self.gen_method_call(val, '__len__', [], int_rprimitive, line)
+            length = self.coerce(length, int_rprimitive, line)
+            ok, fail = BasicBlock(), BasicBlock()
+            self.compare_tagged_condition(length, Integer(0), '>=', ok, fail, line)
+            self.activate_block(fail)
+            self.add(RaiseStandardError(RaiseStandardError.VALUE_ERROR,
+                                        "__len__() should return >= 0",
+                                        line))
+            self.add(Unreachable())
+            self.activate_block(ok)
+            return length
         else:
+            # generic case
             if use_pyssize_t:
                 return self.call_c(generic_ssize_t_len_op, [val], line)
             else:

mypyc/irbuild/specialize.py

@@ -85,7 +85,7 @@ def translate_len(
             return Integer(len(expr_rtype.types))
         else:
             obj = builder.accept(expr.args[0])
-            return builder.builtin_len(obj, -1)
+            return builder.builtin_len(obj, expr.line)
     return None
 
 

mypyc/lib-rt/CPy.h

@@ -540,6 +540,7 @@ int CPyStatics_Initialize(PyObject **statics,
                           const double *floats,
                           const double *complex_numbers,
                           const int *tuples);
+PyObject *CPy_Super(PyObject *builtins, PyObject *self);
 
 #ifdef __cplusplus
 }

mypyc/lib-rt/misc_ops.c

@@ -1,4 +1,4 @@
-// Misc primitive operations
+// Misc primitive operations + C helpers
 //
 // These are registered in mypyc.primitives.misc_ops.
 
@@ -631,3 +631,15 @@ int CPyStatics_Initialize(PyObject **statics,
     }
     return 0;
 }
+
+// Call super(type(self), self)
+PyObject *
+CPy_Super(PyObject *builtins, PyObject *self) {
+    PyObject *super_type = PyObject_GetAttrString(builtins, "super");
+    if (!super_type)
+        return NULL;
+    PyObject *result = PyObject_CallFunctionObjArgs(
+        super_type, (PyObject*)self->ob_type, self, NULL);
+    Py_DECREF(super_type);
+    return result;
+}