Inherit slots instead of redefining them (#3105)

Author: kszucs

ibis/expr/operations/core.py

@@ -1,8 +1,6 @@
 import collections
 import functools
 import itertools
-from typing import Any as AnyT
-from typing import Dict
 
 import toolz
 from public import public
@@ -95,37 +93,6 @@ def _pp(x):
 
         return '{}({})'.format(opname, ', '.join(pprint_args))
 
-    def __getstate__(self) -> Dict[str, AnyT]:
-        """The attributes _expr_cached and _hash are
-        used as caches; they can be excluded from
-        serialization without affecting correctness.
-
-        Excluding _expr_cached and _hash from serialization
-        will allow the serialized bytes to be the same for
-        equivalent Node objets.
-
-        Returns
-        -------
-        Dict[str, Any]
-            A dictionary storing the objects attributes.
-        """
-        excluded_slots = {'_expr_cached', '_hash'}
-        return {
-            slot: getattr(self, slot)
-            for slot in self.__slots__
-            if slot not in excluded_slots
-        }
-
-    def __setstate__(self, state: Dict[str, AnyT]) -> None:
-        """
-        Parameters
-        ----------
-        state: Dict[str, Any]
-            A dictionary storing the objects attributes.
-        """
-        for slot in state:
-            setattr(self, slot, state[slot])
-
     @property
     def inputs(self):
         return tuple(self.args)

ibis/expr/signature.py

@@ -1,20 +1,14 @@
 import copy
 import inspect
-import typing
+from typing import Any, Callable, Dict
 
 import ibis.expr.rules as rlz
 from ibis import util
 
-try:
-    from cytoolz import unique
-except ImportError:
-    from toolz import unique
-
-
 EMPTY = inspect.Parameter.empty  # marker for missing argument
 
 
-class Validator(typing.Callable):
+class Validator(Callable):
     pass
 
 
@@ -112,37 +106,30 @@ def Argument(validator, default=EMPTY):
 
 class AnnotableMeta(type):
     def __new__(metacls, clsname, bases, dct):
-        slots, params = [], {}
-
+        params = {}
         for parent in bases:
-            # inherit parent slots
-            if hasattr(parent, '__slots__'):
-                slots += parent.__slots__
             # inherit from parent signatures
             if hasattr(parent, '__signature__'):
                 params.update(parent.__signature__.parameters)
 
+        slots = list(dct.pop('__slots__', []))
         attribs = {}
         for name, attrib in dct.items():
             if isinstance(attrib, Validator):
                 # so we can set directly
                 params[name] = Parameter(name, validator=attrib)
+                slots.append(name)
             else:
                 attribs[name] = attrib
 
-        # if slots or signature are defined no inheritance happens
-        slots = attribs.get('__slots__', tuple(slots))
-        slots += tuple(params.keys())
-
         # mandatory fields without default values must preceed the optional
         # ones in the function signature, the partial ordering will be kept
         params = sorted(
             params.values(), key=lambda p: p.default is EMPTY, reverse=True
         )
-        signature = attribs.get('__signature__', inspect.Signature(params))
 
-        attribs['__slots__'] = tuple(unique(slots))
-        attribs['__signature__'] = signature
+        attribs['__slots__'] = tuple(slots)
+        attribs['__signature__'] = inspect.Signature(params)
 
         return super().__new__(metacls, clsname, bases, attribs)
 
@@ -162,11 +149,27 @@ def __init__(self, *args, **kwargs):
     def _validate(self):
         pass
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         if self is other:
             return True
         return type(self) == type(other) and self.args == other.args
 
+    def __getstate__(self) -> Dict[str, Any]:
+        return {
+            key: getattr(self, key)
+            for key in self.__signature__.parameters.keys()
+        }
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """
+        Parameters
+        ----------
+        state: Dict[str, Any]
+            A dictionary storing the objects attributes.
+        """
+        for key, value in state.items():
+            setattr(self, key, value)
+
     @property
     def argnames(self):
         return tuple(self.__signature__.parameters.keys())

ibis/tests/expr/test_signature.py

@@ -1,3 +1,4 @@
+import pickle
 from inspect import Signature
 
 import pytest
@@ -28,6 +29,24 @@ def __call__(self, arg, **kwargs):
         return arg
 
 
+class Op(Annotable):
+    __slots__ = ('_cache', '_hash')
+
+
+class ValueOp(Op):
+    arg = Argument(object)
+
+
+class StringOp(ValueOp):
+    arg = Argument(str)
+
+
+class MagicString(StringOp):
+    foo = Argument(str)
+    bar = Argument(bool)
+    baz = Argument(int)
+
+
 @pytest.mark.parametrize('validator', [3, 'coerce'])
 def test_invalid_validator(validator):
     with pytest.raises(TypeError):
@@ -198,32 +217,77 @@ class NoHooves(Farm):
     assert g1.goats == 0
 
 
+def test_copy_default():
+    default = []
+
+    class Op(Annotable):
+        arg = rlz.optional(rlz.instance_of(list), default=default)
+
+    op = Op()
+    assert op.arg is not default
+
+
 def test_slots_are_inherited_and_overridable():
     class Op(Annotable):
         __slots__ = ('_cache',)  # first definition
         arg = Argument(lambda x: x)
 
     class StringOp(Op):
-        arg = Argument(str)  # inherit
+        arg = Argument(str)  # new overridden slot
 
     class StringSplit(StringOp):
-        sep = Argument(str)  # inherit
+        sep = Argument(str)  # new slot
 
     class StringJoin(StringOp):
-        __slots__ = ('_memoize',)  # override
-        sep = Argument(str)
+        __slots__ = ('_memoize',)  # new slot
+        sep = Argument(str)  # new overridden slot
 
     assert Op.__slots__ == ('_cache', 'arg')
-    assert StringOp.__slots__ == ('_cache', 'arg')
-    assert StringSplit.__slots__ == ('_cache', 'arg', 'sep')
-    assert StringJoin.__slots__ == ('_memoize', 'arg', 'sep')
+    assert StringOp.__slots__ == ('arg',)
+    assert StringSplit.__slots__ == ('sep',)
+    assert StringJoin.__slots__ == ('_memoize', 'sep')
 
 
-def test_copy_default():
-    default = []
+def test_multiple_inheritance():
+    # multiple inheritance is allowed only if one of the parents has non-empty
+    # __slots__ definition, otherwise python will raise lay-out conflict
 
     class Op(Annotable):
-        arg = rlz.optional(rlz.instance_of(list), default=default)
+        __slots__ = ('_hash',)
 
-    op = Op()
-    assert op.arg is not default
+    class ValueOp(Annotable):
+        arg = Argument(object)
+
+    class Reduction(ValueOp):
+        _reduction = True
+
+    class UDF(ValueOp):
+        func = Argument(lambda fn, this: fn)
+
+    class UDAF(UDF, Reduction):
+        arity = Argument(int)
+
+    class A(Annotable):
+        a = Argument(int)
+
+    class B(Annotable):
+        b = Argument(int)
+
+    msg = "multiple bases have instance lay-out conflict"
+    with pytest.raises(TypeError, match=msg):
+
+        class AB(A, B):
+            ab = Argument(int)
+
+    assert UDAF.__slots__ == ('arity',)
+    strlen = UDAF(arg=2, func=lambda value: len(str(value)), arity=1)
+    assert strlen.arg == 2
+    assert strlen.arity == 1
+    assert strlen._reduction is True
+
+
+def test_pickling_support():
+    op = MagicString(arg="something", foo="magic", bar=True, baz=8)
+    raw = pickle.dumps(op)
+    loaded = pickle.loads(raw)
+    assert op == loaded