add_slot.py

add_slot.py

""" Classes and metaclasses for easier ``__slots__`` handling.  """

from itertools import tee
import dis

__version__ = "2021.1.6"
__all__ = ("Slots",)


def self_assignemts(method) -> set:
    """Given a method, collect all the attribute names for assignments
    to "self"."""
    # Get the name of the var used to refer to the instance. Usually,
    # this will be "self". It's the first parameter to the
    # __init__(self, ...)  method call. If there are no parameters,
    # just pretend it was "self".
    instance_var = next(iter(method.__code__.co_varnames), "self")
    # We need to scan all the bytecode instructions to see all the times
    # an attribute of "self" got assigned-to. First get the list of
    # instructions.
    instructions = dis.Bytecode(method)
    # Assignments to attributes of "self" are identified by a first
    # LOAD_FAST (with a value of "self") immediately followed by a
    # STORE_ATTR (with a value of the attribute name). So we will need
    # to look at a sequence of pairs through the bytecode. The easiest
    # way to do this is with two iterators.
    i0, i1 = tee(instructions)
    # March the second one ahead by one step.
    next(i1, None)
    names = set()
    # a and b are a pair of bytecode instructions; b follows a.
    for a, b in zip(i0, i1):
        accessing_self = a.argval == instance_var and a.opname == "LOAD_FAST"
        storing_attribute = b.opname == "STORE_ATTR"
        if accessing_self and storing_attribute:
            names.add(b.argval)
    return names


def super_has_dict(cls) -> bool:
    return hasattr(cls, "__slots__") and "__dict__" in cls.__slots__


class SlotsMeta(type):
    def __new__(metacls, name, bases, classdict):

        noslot = bool(classdict.get("noslot"))
        if noslot is True:
            return super().__new__(metacls, name, bases, classdict)

        # Caller may have already provided slots, in which case just
        # retain them and keep going. Note that we make a set() to make
        # it easier to avoid duplicates.
        slots = set(classdict.get("__slots__", ()))

        if "__init__" in classdict:
            slots |= self_assignemts((classdict["__init__"]))

        classdict["__slots__"] = slots
        return super().__new__(metacls, name, bases, classdict)


class Slots(metaclass=SlotsMeta):
    pass


class A:
    def __init__(self, m, n):
        self.m = m
        self.n = n


class B(Slots):
    def __init__(self, o, p):
        self.o = o
        self.p = p


if __name__ == "__main__":
    import time

    # Initialization
    s0 = time.perf_counter()
    for _ in range(10_000_000):
        A(1, 2)
    print(f"Slotless class took {(time.perf_counter()-s0)} to initialize")

    # Attribute access
    a = A(1, 2)
    s1 = time.perf_counter()
    for _ in range(10_000_000):
        a.m
    print(f"Slotless class took {(time.perf_counter()-s1)} to access attribute m")

    # Initialization
    s2 = time.perf_counter()
    for _ in range(10_000_000):
        B(1, 2)
    print(f"Slotted class took {(time.perf_counter()-s2)} to initialize")

    # Attribute access
    b = B(1, 2)
    s3 = time.perf_counter()
    for _ in range(10_000_000):
        b.o

    print(f"Slotless class took {(time.perf_counter()-s3)} to access attribute o")