Rewrite class during creation with supported slots

During the first creation of the loaded class with slots, a new class is
inserted into the inheritance chain to ensure that the slot fields are defined.

This is done because of the need to create an empty struct during the
parsing phase and fill it in later - slots require the fields to be known before hand.

All checks on the new replacement class will have it looking like the
original except an equality comparison between the replaced class and
its replacement.

>>> import thriftpy
>>> ab = thriftpy.load('addressbook.thrift', use_slots=True)
>>> ab_inst = ab.AddressBook()
>>> ab_inst.__class__ == ab.AddressBook # will return False
>>> # all other checks should work as expected
>>> isinstance(ab_inst, ab.AddressBook) # will return True
>>> issubclass(ab_inst.__class__, ab.AddressBook) # will return True

In order to get pickling to work as expected, a new extension type is
registered with copyreg (copy_reg for py2x) to avoid pickling errors.

Author: misakwa

tests/test_hook.py

@@ -82,38 +82,54 @@ def test_load_slots():
     )
 
     # normal structs will have slots
-    # assert set(thrift.PhoneNumber.__slots__) == set(['type', 'number', 'mix_item'])
-    # assert set(thrift.Person.__slots__) == set(['name', 'phones', 'created_at'])
-    # assert set(thrift.AddressBook.__slots__) == set(['people'])
-    # assert set(thrift.AddressBookService.__slots__) == set()
+    assert thrift.PhoneNumber.__slots__ == ['type', 'number', 'mix_item']
+    assert thrift.Person.__slots__ == ['name', 'phones', 'created_at']
+    assert thrift.AddressBook.__slots__ == ['people']
 
-    # one cannot get/set undefined attributes
-    # person = thrift.Person()
-    # with pytest.raises(AttributeError):
-    #     person.attr_not_exist = "Does not work"
+    # get/set undefined attributes
+    person = thrift.Person()
+    with pytest.raises(AttributeError):
+        person.attr_not_exist = "Does not work"
 
-    # with pytest.raises(AttributeError):
-    #     person.attr_not_exist
+    with pytest.raises(AttributeError):
+        person.attr_not_exist
 
-    # pn = thrift.PhoneNumber()
-    # with pytest.raises(AttributeError):
-    #     pn.attr_not_exist = "Does not work"
+    pn = thrift.PhoneNumber()
+    with pytest.raises(AttributeError):
+        pn.attr_not_exist = "Does not work"
 
-    # with pytest.raises(AttributeError):
-    #     pn.attr_not_exist
+    with pytest.raises(AttributeError):
+        pn.attr_not_exist
 
-    # ab = thrift.AddressBook()
-    # with pytest.raises(AttributeError):
-    #     ab.attr_not_exist = "Does not work"
+    ab = thrift.AddressBook()
+    with pytest.raises(AttributeError):
+        ab.attr_not_exist = "Does not work"
 
-    # with pytest.raises(AttributeError):
-    #     ab.attr_not_exist
+    with pytest.raises(AttributeError):
+        ab.attr_not_exist
+    # eo: get/set
 
     # exceptions will not have slots
-    # assert not hasattr(thrift.PersonNotExistsError, '__slots__')
+    assert not hasattr(thrift.PersonNotExistsError, '__slots__')
 
     # enums will not have slots
-    # assert not hasattr(thrift.PhoneType, '__slots__')
+    assert not hasattr(thrift.PhoneType, '__slots__')
+
+    # service itself will not be created with slots
+    assert not hasattr(thrift.AddressBookService, '__slots__')
+
+    # service args will have slots
+    args_slots = thrift.AddressBookService.get_phonenumbers_args.__slots__
+    assert args_slots == ['name', 'count']
+
+    # XXX: service result will have their slot list empty until after they are
+    # created. This is because the success field is inserted after calling
+    # _make_struct. We could hardcode a check in the metaclass for names ending
+    # with _result, but I'm unsure if its a good idea. This should usually not
+    # be an issue since this object is only used internally, but we can revisit
+    # the need to have it available when required.
+    result_obj = thrift.AddressBookService.get_phonenumbers_result()
+    assert result_obj.__slots__ == ['success']
 
     # should be able to pickle slotted objects - if load with module_name
     bob = thrift.Person(name="Bob")

thriftpy/parser/parser.py

@@ -767,6 +767,8 @@ def _make_enum(name, kvs):
 
 def _make_empty_struct(name, ttype=TType.STRUCT, base_cls=TPayload):
     attrs = {'__module__': thrift_stack[-1].__name__, '_ttype': ttype}
+    if issubclass(base_cls, TSPayload):
+        attrs['__slots__'] = []
     return type(name, (base_cls, ), attrs)
 
 
@@ -787,10 +789,9 @@ def _fill_in_struct(cls, fields, _gen_init=True):
     setattr(cls, 'thrift_spec', thrift_spec)
     setattr(cls, 'default_spec', default_spec)
     setattr(cls, '_tspec', _tspec)
-    # add __slots__ so we can check way before the class is created, even though
-    # it really does nothing here, the real work is done during instantiation
+    # add __slots__ for easy introspection
     if issubclass(cls, TSPayload):
-        cls.__slots__ = tuple(field for field, _ in default_spec)
+        cls.__slots__ = [field for field, _ in default_spec]
     if _gen_init:
         gen_init(cls, thrift_spec, default_spec)
     return cls
@@ -808,8 +809,7 @@ def _make_service(name, funcs, extends, use_slots=False):
 
     attrs = {'__module__': thrift_stack[-1].__name__}
     base_cls = TSPayload if use_slots else TPayload
-    if use_slots:
-        attrs['__slots__'] = tuple()
+    # service class itself will not be created with slots
     cls = type(name, (extends, ), attrs)
     thrift_services = []
 

thriftpy/thrift.py

@@ -9,10 +9,13 @@
 
 from __future__ import absolute_import
 
+try:
+    import copy_reg as copyreg
+except ImportError:
+    import copyreg
+
 import functools
-import itertools
 import linecache
-import operator
 import types
 
 from ._compat import with_metaclass
@@ -136,57 +139,27 @@ def __new__(cls, name, bases, attrs):
             attrs["__init__"] = init_func_generator(cls, spec)
         return super(TPayloadMeta, cls).__new__(cls, name, bases, attrs)
 
-# checks: instance and subclass checks
-    def __instancecheck__(cls, inst):
-        mod = inst.__class__.__module__
-        name = inst.__class__.__name__
-        key = "%s:%s" % (mod, name)
-        repl_cls = TPayloadMeta._class_cache.get(key)
-        if repl_cls is None:
-            return type.__instancecheck__(cls, inst)
-        return inst.__class__ is repl_cls
-
-    def __subclasscheck__(cls, subcls):
-        mod = cls.__module__
-        name = cls.__name__
-        key = "%s:%s" % (mod, name)
-        repl_cls = TPayloadMeta._class_cache.get(key)
-        if repl_cls is None:
-            return type.__subclasscheck__(cls, subcls)
-        if cls == repl_cls:
-            return True
-        # the first class in __mro__ is replaced by the replacement class so we
-        # can only look up from the second position
-        return cls.__mro__[1] in repl_cls.__mro__
-# eo: checks
-
     def __call__(cls, *args, **kw):
-        if not cls.__mro__[1] == TSPayload:
-            return super(TPayloadMeta, cls).__call__(cls, *args, **kw)
-        # XXX: replaces class with new class using slot list from default_spec
+        # if issubclass(cls, TSPayload):
+        if not issubclass(cls, TSPayload):
+            return type.__call__(cls, *args, **kw)
         cls_name = cls.__name__.split('.')[-1]
         cache_key = '%s:%s' % (cls.__module__, cls_name)
-        # XXX: we trust default_spec more than __slots__ for this
-        fields = tuple(field for field, _ in cls.default_spec)
-        cls_obj = TPayloadMeta._class_cache.get(cache_key)
-        if not cls_obj:
-            cls_obj = type(
+        kls = TPayloadMeta._class_cache.get(cache_key)
+        if not kls:
+            fields = [field for field, _ in cls.default_spec]
+            kls = type(
                 cls_name,
-                cls.__mro__,
+                (cls,),
                 {
                     '__slots__': fields,
                     '__module__': cls.__module__,
                 }
             )
-            # XXX: need a better way to do this; its a dupe from parser.py
-            cls_obj._ttype = cls._ttype
-            cls_obj._tspec = cls._tspec
-            cls_obj.default_spec = cls.default_spec
-            cls_obj.thrift_spec = cls.thrift_spec
-            # cls.__init__ is already bound to cls
-            cls_obj.__init__ = init_func_generator(cls_obj, cls.default_spec)
-            TPayloadMeta._class_cache[cache_key] = cls_obj
-        return type.__call__(cls_obj, *args, **kw)
+            TPayloadMeta._class_cache[cache_key] = kls
+            fn = lambda obj: (cls, tuple(getattr(obj, f) for f in fields))
+            copyreg.pickle(kls, fn)
+        return type.__call__(kls, *args, **kw)
 
 
 def gen_init(cls, thrift_spec=None, default_spec=None):
@@ -236,11 +209,8 @@ def write(self, oprot):
         oprot.write_struct(self)
 
     def __repr__(self):
-        keys = itertools.chain.from_iterable(
-            getattr(cls, '__slots__', tuple()) for cls in type(self).__mro__
-        )
-        keys = list(keys)
-        values = operator.attrgetter(*keys)(self)
+        keys = self.__slots__
+        values = [getattr(self, k) for k in keys]
         l = ['%s=%r' % (key, value) for key, value in zip(keys, values)]
         return '%s(%s)' % (self.__class__.__name__, ', '.join(l))
 
@@ -250,28 +220,14 @@ def __str__(self):
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return False
-        keys = itertools.chain.from_iterable(
-            getattr(cls, '__slots__', tuple()) for cls in type(self).__mro__
-        )
-        keys = list(keys)
-        getter = operator.attrgetter(*keys)
-        return getter(self) == getter(other)
+        keys = self.__slots__
+        vals1 = [getattr(self, k) for k in keys]
+        vals2 = [getattr(self, k) for k in keys]
+        return vals1 == vals2
 
     def __ne__(self, other):
         return not self.__eq__(other)
 
-    def __getstate__(self):
-        keys = itertools.chain.from_iterable(
-            getattr(cls, '__slots__', tuple()) for cls in type(self).__mro__
-        )
-        keys = list(keys)
-        values = operator.attrgetter(*keys)(self)
-        return tuple(zip(keys, values))
-
-    def __setstate__(self, state):
-        for k, v in state:
-            setattr(self, k, v)
-
 
 class TClient(object):