deslot.py

"""
In [1]: class C:
   ...:     __slots__ = 'a', 'b'
   ...:

In [2]: C.__slots__
Out[2]: ('a', 'b')

In [3]: from deslot import deslot

In [4]: with deslot():
   ...:     class D:
   ...:         __slots__ = 'a', 'b'
   ...:

In [5]: d = D()

In [6]: d.c = 'ayy'

In [7]: d.c
Out[7]: 'ayy

In [8]: D.__slots__
Out[8]: ('a', 'b')

In [9]: from deslot import autoslot

In [10]: with autoslot():
    ...:     class E:
    ...:         def __init__(self, a, b):
    ...:             self.a = a
    ...:             self.b = b
    ...:             self.c = a + b
    ...:

In [11]: e = E(1, 2)

In [12]: E.__slots__
Out[12]: ('a', 'b', 'c')

In [13]: e.d = 1
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-13-9977a258af52> in <module>()
----> 1 e.d = 1

AttributeError: 'E' object has no attribute 'd'
"""
import builtins
import contextlib

from codetransformer import Code, instructions
import toolz

_real_build_class = builtins.__build_class__


def _deslot_build_class(func, name, *bases, metaclass=None, **kwargs):
    metaclass = metaclass if metaclass is not None else type

    slots = no_slots = object()

    def new(name, bases, dict_, **kwargs):
        nonlocal slots
        slots = dict_.pop('__slots__', no_slots)
        return metaclass.__new__(metaclass, name, bases, dict_, **kwargs)

    cls = _real_build_class(func, name, *bases, metaclass=new, **kwargs)
    if slots is not no_slots:
        cls.__slots__ = slots
    return cls


@contextlib.contextmanager
def deslot():
    builtins.__build_class__ = _deslot_build_class
    try:
        yield
    finally:
        builtins.__build_class__ = _real_build_class


def _autoslot_build_class(func, name, *bases, metaclass=None, **kwargs):
    metaclass = metaclass if metaclass is not None else type

    def new(name, bases, dict_, **kwargs):
        ghost_class = metaclass.__new__(
            metaclass,
            name,
            bases,
            dict_,
            **kwargs,
        )
        existing_slots = getattr(ghost_class, '__slots__', set())
        init_instrs = Code.from_pyfunc(ghost_class.__init__)
        slots = {
            store.arg
            for ob, store in toolz.sliding_window(2, init_instrs)
            if isinstance(ob, instructions.LOAD_FAST) and
            ob.arg == 'self' and
            isinstance(store, instructions.STORE_ATTR)
        }
        dict_['__slots__'] = tuple(sorted(slots - existing_slots))
        return metaclass.__new__(metaclass, name, bases, dict_, **kwargs)

    return _real_build_class(func, name, *bases, metaclass=new, **kwargs)


@contextlib.contextmanager
def autoslot():
    builtins.__build_class__ = _autoslot_build_class
    try:
        yield
    finally:
        builtins.__build_class__ = _real_build_class