# /// script
# dependencies = ["numpy"]
# ///
"""How many bytes does one small Python object cost?

Creates N instances of a few variants of the same 3-attribute class and
measures the memory they add with tracemalloc, plus how fast their
attributes read. Then wide classes, which fall off CPython's fast path, and
no objects at all: the attributes as raw ints in NumPy arrays.

    uv run --python 3.13 measure.py

Written for "A million objects: where do the bytes go?":
https://ksh.fyi/blogs/python-slots/
"""
import gc
import sys
import time
import tracemalloc

N = 1_000_000


class Tree:
    def __init__(self, x, y, age):
        self.x = x
        self.y = y
        self.age = age


class TreeOwnDict:
    """Opts out of key-sharing: every instance gets its own ordinary dict,
    roughly how all instances looked before Python 3.3 (PEP 412)."""

    def __init__(self, x, y, age):
        self.__dict__ = {"x": x, "y": y, "age": age}


class TreeSlots:
    __slots__ = ("x", "y", "age")

    def __init__(self, x, y, age):
        self.x = x
        self.y = y
        self.age = age


def make(cls):
    objs = [None] * N  # allocate the list up front so it isn't counted
    gc.collect()
    tracemalloc.start()
    before = tracemalloc.get_traced_memory()[0]
    for i in range(N):
        # small ints (< 256) are cached by CPython, so they cost nothing here
        objs[i] = cls(i % 100, i % 200, i % 50)
    used = tracemalloc.get_traced_memory()[0] - before
    tracemalloc.stop()
    return objs, used / N


def touch_dict(objs):
    """Reading __dict__ forces 3.11+ to build a real dict object."""
    gc.collect()
    tracemalloc.start()
    before = tracemalloc.get_traced_memory()[0]
    for o in objs:
        o.__dict__
    used = tracemalloc.get_traced_memory()[0] - before
    tracemalloc.stop()
    return used / N


def read_ns(objs):
    start = time.perf_counter()
    for o in objs:
        o.x
        o.y
        o.age
    return (time.perf_counter() - start) / (3 * N) * 1e9


def age_ns(objs):
    """Age every tree by a year, one object at a time."""
    start = time.perf_counter()
    for o in objs:
        o.age += 1
    return (time.perf_counter() - start) / N * 1e9


def numpy_forest():
    """No Tree objects at all: one NumPy array of raw 8-byte ints per attribute."""
    import numpy as np

    gc.collect()
    tracemalloc.start()  # NumPy reports its array buffers to tracemalloc
    before = tracemalloc.get_traced_memory()[0]
    i = np.arange(N)
    x, y, age = i % 100, i % 200, i % 50
    del i
    used = tracemalloc.get_traced_memory()[0] - before
    tracemalloc.stop()

    start = time.perf_counter()
    age += 1
    aging = (time.perf_counter() - start) / N * 1e9

    start = time.perf_counter()
    for k in range(N):
        x[k]
        y[k]
        age[k]
    read = (time.perf_counter() - start) / (3 * N) * 1e9
    return used / N, aging, read


def wide(n, slots=False):
    """A class that sets n attributes in __init__."""
    names = ["a%d" % k for k in range(n)]

    def __init__(self, *args):
        for name in names:
            setattr(self, name, 1)

    ns = {"__init__": __init__}
    if slots:
        ns["__slots__"] = tuple(names)
    return type("Wide%d" % n, (), ns)


def main():
    version = "%d.%d" % sys.version_info[:2]
    rows = []
    aging = {}
    for cls in (TreeOwnDict, Tree, TreeSlots):
        objs, per_obj = make(cls)
        rows.append((cls.__name__, per_obj, sys.getsizeof(objs[0]), read_ns(objs)))
        aging[cls.__name__] = age_ns(objs)
        if cls is Tree:
            extra = touch_dict(objs)
            rows.append(("Tree, __dict__ read", per_obj + extra,
                         sys.getsizeof(objs[0]), read_ns(objs)))
        del objs

    print("Python %s, %s instances" % (version, format(N, ",")))
    print("%-22s %10s %10s %10s" % ("variant", "bytes/obj", "getsizeof", "ns/read"))
    for name, per_obj, size, ns in rows:
        print("%-22s %10.1f %10d %10.1f" % (name, per_obj, size, ns))

    print()
    print("%-22s %10s" % ("wide classes", "bytes/obj"))
    shapes = [("%d attrs" % n, wide(n)) for n in (29, 30)]
    shapes += [("30 attrs, slots", wide(30, slots=True))]
    for name, cls in shapes:
        objs, per_obj = make(cls)
        print("%-22s %10.1f" % (name, per_obj))
        del objs

    print()
    per_tree, numpy_aging, numpy_read = numpy_forest()
    print("NumPy columns: %.1f bytes/tree, %.1f ns/read" % (per_tree, numpy_read))

    print()
    print("%-22s %10s" % ("age every tree", "ns/tree"))
    for name, ns in aging.items():
        print("%-22s %10.2f" % (name + " loop", ns))
    print("%-22s %10.2f" % ("NumPy age += 1", numpy_aging))

if __name__ == "__main__":
    main()