Support defining __bool__ (mypyc/mypyc#373)

msullivan · web-flow · commit 2d761d0bcfa7 · 2018-08-13T13:26:41.000-07:00
To support this, refactor emitclass dunder handling to be even *more*
data driven.
The main reason I am not throwing in support for ~all of the other
dunder methods is that I think we should add tests for them when we
add support.

Supporting __bool__ without always deoptimizing truthiness checks
against Optional types requires adding a very simple "whole-program
analysis" for whether a subclass defines __bool__.
diff --git a/mypyc/emitclass.py b/mypyc/emitclass.py
@@ -10,6 +10,7 @@
 from mypyc.emitfunc import native_function_header
 from mypyc.emitwrapper import (
     generate_dunder_wrapper, generate_hash_wrapper, generate_richcompare_wrapper,
+    generate_bool_wrapper,
 )
 from mypyc.ops import (
     ClassIR, FuncIR, FuncDecl, RType, RTuple, Environment, object_rprimitive, FuncSignature,
@@ -31,6 +32,7 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
 # We maintain a table from dunder function names to struct slots they
 # correspond to and functions that generate a wrapper (if necessary)
 # and return the function name to stick in the slot.
+# TODO: Add remaining dunder methods
 SlotGenerator = Callable[[ClassIR, FuncIR, Emitter], str]
 SlotTable = Mapping[str, Tuple[str, SlotGenerator]]
 
@@ -44,6 +46,19 @@ def wrapper_slot(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     '__hash__': ('tp_hash', generate_hash_wrapper),
 }
 
+AS_MAPPING_SLOT_DEFS = {
+    '__getitem__': ('mp_subscript', generate_dunder_wrapper),
+}
+
+AS_NUMBER_SLOT_DEFS = {
+    '__bool__': ('nb_bool', generate_bool_wrapper),
+}
+
+SIDE_TABLES = [
+    ('as_mapping', 'PyMappingMethods', AS_MAPPING_SLOT_DEFS),
+    ('as_number', 'PyNumberMethods', AS_NUMBER_SLOT_DEFS),
+]
+
 
 def generate_slots(cl: ClassIR, table: SlotTable, emitter: Emitter) -> Dict[str, str]:
     fields = OrderedDict()  # type: Dict[str, str]
@@ -103,12 +118,14 @@ def emit_line() -> None:
     init_fn = cl.get_method('__init__')
 
     # Fill out slots in the type object from dunder methods.
-    # TODO: Add remaining dunder methods
     fields.update(generate_slots(cl, SLOT_DEFS, emitter))
 
-    as_mapping_name = generate_as_mapping_for_class(cl, emitter)
-    if as_mapping_name:
-        fields['tp_as_mapping'] = '&{}'.format(as_mapping_name)
+    # Fill out dunder methods that live in tables hanging off the side.
+    for table_name, type, slot_defs in SIDE_TABLES:
+        slots = generate_slots(cl, slot_defs, emitter)
+        if slots:
+            table_struct_name = generate_side_table_for_class(cl, table_name, type, slots, emitter)
+            fields['tp_{}'.format(table_name)] = '&{}'.format(table_struct_name)
 
     richcompare_name = generate_richcompare_wrapper(cl, emitter)
     if richcompare_name:
@@ -457,18 +474,13 @@ def generate_methods_table(cl: ClassIR,
     emitter.emit_line('};')
 
 
-AS_MAPPING_SLOT_DEFS = {
-    '__getitem__': ('mp_subscript', generate_dunder_wrapper),
-}
-
-
-def generate_as_mapping_for_class(cl: ClassIR,
+def generate_side_table_for_class(cl: ClassIR,
+                                  name: str,
+                                  type: str,
+                                  slots: Dict[str, str],
                                   emitter: Emitter) -> Optional[str]:
-    slots = generate_slots(cl, AS_MAPPING_SLOT_DEFS, emitter)
-    if not slots:
-        return None
-    name = '{}_as_mapping'.format(cl.name_prefix(emitter.names))
-    emitter.emit_line('static PyMappingMethods {} = {{'.format(name))
+    name = '{}_{}'.format(cl.name_prefix(emitter.names), name)
+    emitter.emit_line('static {} {} = {{'.format(type, name))
     for field, value in slots.items():
         emitter.emit_line(".{} = {},".format(field, value))
     emitter.emit_line("};")
diff --git a/mypyc/emitwrapper.py b/mypyc/emitwrapper.py
@@ -4,7 +4,8 @@
 from mypyc.emit import Emitter
 from mypyc.ops import (
     ClassIR, FuncIR, RType, RuntimeArg,
-    is_object_rprimitive, is_int_rprimitive,
+    is_object_rprimitive, is_int_rprimitive, is_bool_rprimitive,
+    bool_rprimitive,
     FUNC_STATICMETHOD,
 )
 from mypyc.namegen import NameGenerator
@@ -131,6 +132,26 @@ def generate_hash_wrapper(cl: ClassIR, fn: FuncIR, emitter: Emitter) -> str:
     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))
+    emitter.emit_line('static int {name}(PyObject *self) {{'.format(
+        name=name
+    ))
+    emitter.emit_line('{}val = {}{}(self);'.format(emitter.ctype_spaced(fn.ret_type),
+                                                   NATIVE_PREFIX,
+                                                   fn.cname(emitter.names)))
+    emitter.emit_error_check('val', fn.ret_type, 'return -1;')
+    # This wouldn't be that hard to fix but it seems unimportant and
+    # getting error handling and unboxing right would be fiddly. (And
+    # way easier to do in IR!)
+    assert is_bool_rprimitive(fn.ret_type), "Only bool return supported for __bool__"
+    emitter.emit_line('return val;')
+    emitter.emit_line('}')
+
+    return name
+
+
 def generate_wrapper_core(fn: FuncIR, emitter: Emitter,
                           optional_args: List[RuntimeArg] = [],
                           arg_names: Optional[List[str]] = None) -> None:
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -380,6 +380,12 @@ def prepare_class_def(module_name: str, cdef: ClassDef, mapper: Mapper) -> None:
     ir.mro = mro
     ir.base_mro = base_mro
 
+    # We need to know whether any children of a class have a __bool__
+    # method in order to know whether we can assume it is always true.
+    if ir.has_method('__bool__'):
+        for base in ir.mro:
+            base.has_bool = True
+
 
 class FuncInfo(object):
     """Contains information about functions as they are generated."""
@@ -2329,9 +2335,8 @@ def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) ->
                 is_none = self.binary_op(value, self.none(), 'is not', value.line)
                 branch = Branch(is_none, true, false, Branch.BOOL_EXPR)
                 self.add(branch)
-                if isinstance(value_type, RInstance):
+                if isinstance(value_type, RInstance) and not value_type.class_ir.has_bool:
                     # Optional[X] where X is always truthy
-                    # TODO: Support __bool__
                     pass
                 else:
                     # Optional[X] where X may be falsey and requires a check
diff --git a/mypyc/ops.py b/mypyc/ops.py
@@ -1431,6 +1431,9 @@ def __init__(self, name: str, module_name: str, is_trait: bool = False,
         # base_mro is the chain of concrete (non-trait) ancestors
         self.base_mro = [self]  # type: List[ClassIR]
 
+        # Does this class or any subclass have a __bool__method
+        self.has_bool = False
+
     def real_base(self) -> Optional['ClassIR']:
         """Return the actual concrete base class, if there is one."""
         if len(self.mro) > 1 and not self.mro[1].is_trait:
diff --git a/test-data/fixtures/ir.py b/test-data/fixtures/ir.py
@@ -64,7 +64,9 @@ def __eq__(self, x:object) -> bool:pass
     def __ne__(self, x: object) -> bool: pass
     def join(self, x: Iterable[object]) -> bytes: pass
 
-class bool: pass
+class bool:
+    def __init__(self, o: object = ...) -> None: ...
+
 
 class tuple(Generic[T], Sized):
     def __init__(self, i: Iterable[T]) -> None: pass
diff --git a/test-data/genops-optional.test b/test-data/genops-optional.test
@@ -24,6 +24,49 @@ L2:
     r3 = 2
     return r3
 
+[case testIsTruthyOverride]
+from typing import Optional
+
+class A: pass
+
+class B(A):
+    def __bool__(self) -> bool:
+        return False
+
+
+def f(x: Optional[A]) -> int:
+    if x:
+        return 1
+    return 2
+[out]
+def B.__bool__(self):
+    self :: B
+    r0 :: bool
+L0:
+    r0 = False
+    return r0
+def f(x):
+    x :: union[A, None]
+    r0 :: None
+    r1 :: bool
+    r2 :: A
+    r3 :: bool
+    r4, r5 :: int
+L0:
+    r0 = None
+    r1 = x is not r0
+    if r1 goto L1 else goto L3 :: bool
+L1:
+    r2 = cast(A, x)
+    r3 = bool r2 :: object
+    if r3 goto L2 else goto L3 :: bool
+L2:
+    r4 = 1
+    return r4
+L3:
+    r5 = 2
+    return r5
+
 [case testIsNotNone]
 from typing import Optional
 
diff --git a/test-data/run.test b/test-data/run.test
@@ -2377,7 +2377,8 @@ class Item:
         return self.value < x.value
 
 class Subclass1(Item):
-    pass
+    def __bool__(self) -> bool:
+        return bool(self.value)
 
 class NonBoxedThing:
     def __getitem__(self, index: Item) -> Item:
@@ -2400,6 +2401,9 @@ def internal_index_into() -> None:
     z = Item("3")
     print(y[z].value)
 
+def is_truthy(x: Item) -> bool:
+    return True if x else False
+
 [file driver.py]
 from native import *
 x = BoxedThing()
@@ -2427,6 +2431,11 @@ assert i2 < i3
 assert not i1 < i2
 assert i1 == Subclass1('lolol')
 
+assert is_truthy(Item(''))
+assert is_truthy(Item('a'))
+assert not is_truthy(Subclass1(''))
+assert is_truthy(Subclass1('a'))
+
 internal_index_into()
 [out]
 7 7
