__slots__でクラスのメモリ使用量を節約する方法(Python)

# Python 3.8+ を想定

class Point:
    # __slots__ を使わない通常クラス
    def __init__(self, x, y):
        self.x = x
        self.y = y

p = Point(1, 2)
print(p.__dict__)  # 属性テーブル(辞書)が存在する

class SlottedPoint:
    # インスタンスが持てる属性名を固定
    __slots__ = ('x', 'y')

    def __init__(self, x, y):
        # 各属性はスロットに直接格納される
        self.x = x
        self.y = y

sp = SlottedPoint(1, 2)

# __dict__ は存在しない(動的属性追加をしないための証)
print(hasattr(sp, '__dict__'))
# 許可された属性へのアクセスは通常通り
print(sp.x, sp.y)

class User:
    __slots__ = ('name', 'age', 'active')

    def __init__(self, name='unknown', age=0, active=True):
        # ここでインスタンス属性の初期値を設定する
        self.name = name
        self.age = age
        self.active = active

u1 = User()                 # すべてデフォルト
u2 = User('Alice', 30)      # 一部だけ指定

print(u1.name, u1.age, u1.active)
print(u2.name, u2.age, u2.active)

class Account:
    __slots__ = ('id', 'balance')

    def __init__(self, id, balance=0):
        self.id = id
        self.balance = balance

acc = Account('A-001', 100)

# 許可されていない属性を追加しようとするとエラー
try:
    acc.owner = 'Bob'  # owner は __slots__ にない
except AttributeError as e:
    print('AttributeError:', e)

import sys
import tracemalloc

# 通常クラス
class Normal:
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

# __slots__ クラス
class Slotted:
    __slots__ = ('a', 'b', 'c')
    def __init__(self, a, b, c):
        self.a = a
        self.b = b
        self.c = c

def measure_allocation(make_objects, n=100_000):
    # n 個のインスタンスを作って増分メモリを測る
    tracemalloc.start()
    snap1 = tracemalloc.take_snapshot()
    objs = make_objects(n)
    snap2 = tracemalloc.take_snapshot()
    stats = snap2.compare_to(snap1, 'lineno')
    total = sum(s.size_diff for s in stats)  # 増分合計(bytes)
    # 生存させ続けないと解放されてしまうので戻り値で保持
    return total, objs

def make_normal(n):
    return [Normal(i, i+1, i+2) for i in range(n)]

def make_slotted(n):
    return [Slotted(i, i+1, i+2) for i in range(n)]

# 参考: getsizeof は"そのもの"のサイズのみ
tmp_n = Normal(1,2,3)
tmp_s = Slotted(1,2,3)
print('getsizeof(Normal instance) =', sys.getsizeof(tmp_n))
print('getsizeof(Slotted instance) =', sys.getsizeof(tmp_s))
print('has __dict__ ?', hasattr(tmp_n, '__dict__'), hasattr(tmp_s, '__dict__'))

# 大量生成時の増分メモリ
normal_bytes, normals = measure_allocation(make_normal, n=100_000)
slotted_bytes, slotteds = measure_allocation(make_slotted, n=100_000)

print('Normal total bytes:', normal_bytes)
print('Slotted total bytes:', slotted_bytes)
print('Per-instance diff (approx):', (normal_bytes - slotted_bytes) // 100_000, 'bytes')

class Base:
    __slots__ = ('x',)

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

# 子が __slots__ を定義しない → __dict__ が復活
class ChildLoosen(Base):
    pass

# 子が属性追加しない場合の推奨形
class ChildTight(Base):
    __slots__ = ()  # 新規スロット無し。__dict__ も作らない

# 子が属性を追加したい場合
class ChildAdd(Base):
    __slots__ = ('y',)  # 親の 'x' に加えて 'y' を追加

a = ChildLoosen(1)
b = ChildTight(1)
c = ChildAdd(1)
c.y = 10

print('Loosen has __dict__:', hasattr(a, '__dict__'))
print('Tight has __dict__:', hasattr(b, '__dict__'))
print('Add has __dict__  :', hasattr(c, '__dict__'))
print('c.x, c.y =', c.x, c.y)

import weakref

class Node:
    __slots__ = ('value', '__weakref__')  # 弱参照を許可
    def __init__(self, value):
        self.value = value

n = Node(10)
r = weakref.ref(n)   # 弱参照が作れる
print(r() is n)

class Extensible:
    __slots__ = ('a', '__dict__')  # 動的属性を許可(→メモリ節約は減る)
    def __init__(self, a):
        self.a = a

e = Extensible(1)
e.b = 2  # __dict__ があるので追加できる
print(e.a, e.b)