Claude Code for tracemalloc: Python Memory Allocation Tracing

Python’s tracemalloc module traces memory block allocations, enabling identification of memory leaks and hotspots. import tracemalloc. start: tracemalloc.start(25) — begin tracing; 25 = frames in call stack per allocation. stop: tracemalloc.stop(). take_snapshot: snap = tracemalloc.take_snapshot() → Snapshot object. get_traced_memory: cur, peak = tracemalloc.get_traced_memory() → bytes. get_tracemalloc_memory: overhead of tracemalloc itself in bytes. clear_traces: tracemalloc.clear_traces() — reset counts. tracemalloc.is_tracing(). Snapshot.statistics: stats = snap.statistics("lineno") — group by "lineno", "filename", "traceback"; returns list of Statistic sorted by size descending. Snapshot.compare_to: diffs = snap2.compare_to(snap1, "lineno") → list of StatisticDiff (size_diff, count_diff). filter_traces: snap.filter_traces([tracemalloc.Filter(inclusive, pattern)]) — include/exclude tracebacks. Statistic: .traceback, .size, .count. StatisticDiff: .size_diff, .count_diff, .size, .count. Traceback: .format(limit=10). tracemalloc.get_object_traceback(obj) — traceback for a specific object. Filter: tracemalloc.Filter(True, "mymodule*") — inclusive filter. Claude Code generates leak detectors, allocation reporters, before/after diff tools, and memory budget guards.

CLAUDE.md for tracemalloc

## tracemalloc Stack
- Stdlib: import tracemalloc
- Start:  tracemalloc.start(25)   # 25 = frames depth
- Snap:   snap = tracemalloc.take_snapshot()
- Top:    for s in snap.statistics("lineno")[:10]: print(s)
- Diff:   diffs = snap2.compare_to(snap1, "lineno")
- Peak:   cur, peak = tracemalloc.get_traced_memory()
- Stop:   tracemalloc.stop()

tracemalloc Memory Tracing Pipeline

# app/tracemallocutil.py — trace, snapshot, top, diff, leak monitor, budget
from __future__ import annotations

import gc
import linecache
import tracemalloc
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Callable, Generator


# ─────────────────────────────────────────────────────────────────────────────
# 1. Tracing context manager
# ─────────────────────────────────────────────────────────────────────────────

@contextmanager
def trace_memory(
    nframe: int = 25,
    gc_collect: bool = True,
) -> Generator[None, None, None]:
    """
    Context manager that enables tracemalloc for the enclosed block.

    nframe: number of call stack frames captured per allocation.
    gc_collect: force GC before and after to reduce noise.

    Example:
        with trace_memory():
            my_function()
        snap = tracemalloc.take_snapshot()
        print_top(snap, limit=10)
    """
    if gc_collect:
        gc.collect()
    tracemalloc.start(nframe)
    try:
        yield
    finally:
        if gc_collect:
            gc.collect()
        tracemalloc.stop()


# ─────────────────────────────────────────────────────────────────────────────
# 2. Snapshot helpers
# ─────────────────────────────────────────────────────────────────────────────

def take_snapshot(
    exclude_stdlib: bool = True,
    exclude_tracemalloc: bool = True,
) -> tracemalloc.Snapshot:
    """
    Take a filtered memory snapshot.

    Example:
        snap = take_snapshot()
        print_top(snap, limit=10)
    """
    import sys
    snap = tracemalloc.take_snapshot()
    filters: list[tracemalloc.Filter] = []
    if exclude_tracemalloc:
        filters.append(tracemalloc.Filter(False, tracemalloc.__file__))
    if exclude_stdlib:
        import sysconfig
        stdlib_path = sysconfig.get_path("stdlib")
        if stdlib_path:
            filters.append(tracemalloc.Filter(False, stdlib_path + "/*"))
    if filters:
        snap = snap.filter_traces(filters)
    return snap


@dataclass
class MemoryStat:
    file:    str
    line:    int
    size:    int       # bytes
    count:   int
    code:    str       # source line text

    def __str__(self) -> str:
        size_kb = self.size / 1024
        return (f"{size_kb:8.1f} KiB  "
                f"{self.count:6d} objects  "
                f"{self.file}:{self.line}  {self.code.strip()!r}")


def top_allocations(
    snap: tracemalloc.Snapshot,
    limit: int = 20,
    key: str = "lineno",
) -> list[MemoryStat]:
    """
    Return the top allocation sites from a snapshot.
    key: "lineno" (default), "filename", "traceback".

    Example:
        with trace_memory():
            workload()
        snap = take_snapshot()
        for s in top_allocations(snap, limit=10):
            print(s)
    """
    stats = snap.statistics(key)
    results: list[MemoryStat] = []
    for stat in stats[:limit]:
        tb = stat.traceback
        frame = tb[0] if tb else None
        file = frame.filename if frame else "?"
        line = frame.lineno if frame else 0
        code = linecache.getline(file, line) if frame else ""
        results.append(MemoryStat(
            file=file, line=line, size=stat.size, count=stat.count, code=code,
        ))
    return results


def print_top(snap: tracemalloc.Snapshot, limit: int = 10) -> None:
    """Print top allocation sites to stdout."""
    for s in top_allocations(snap, limit=limit):
        print(f"  {s}")


# ─────────────────────────────────────────────────────────────────────────────
# 3. Diff / leak detection
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class AllocDiff:
    file:      str
    line:      int
    size_diff: int     # bytes (positive = grew)
    count_diff: int
    size_now:  int
    code:      str

    @property
    def is_growth(self) -> bool:
        return self.size_diff > 0

    def __str__(self) -> str:
        sign = "+" if self.size_diff >= 0 else ""
        return (f"{sign}{self.size_diff / 1024:+.1f} KiB  "
                f"{sign}{self.count_diff:+d} objects  "
                f"{self.file}:{self.line}  {self.code.strip()!r}")


def diff_snapshots(
    before: tracemalloc.Snapshot,
    after: tracemalloc.Snapshot,
    limit: int = 20,
    key: str = "lineno",
    only_growth: bool = True,
) -> list[AllocDiff]:
    """
    Diff two snapshots and return sorted allocation deltas.

    Example:
        snap1 = take_snapshot()
        do_work()
        snap2 = take_snapshot()
        for d in diff_snapshots(snap1, snap2):
            print(d)
    """
    stats = after.compare_to(before, key)
    if only_growth:
        stats = [s for s in stats if s.size_diff > 0]
    results: list[AllocDiff] = []
    for stat in stats[:limit]:
        tb = stat.traceback
        frame = tb[0] if tb else None
        file = frame.filename if frame else "?"
        line = frame.lineno if frame else 0
        code = linecache.getline(file, line) if frame else ""
        results.append(AllocDiff(
            file=file, line=line,
            size_diff=stat.size_diff, count_diff=stat.count_diff,
            size_now=stat.size, code=code,
        ))
    results.sort(key=lambda d: -d.size_diff)
    return results


# ─────────────────────────────────────────────────────────────────────────────
# 4. Memory budget / guard
# ─────────────────────────────────────────────────────────────────────────────

class MemoryBudget:
    """
    Context manager that asserts memory growth stays within a budget.
    Raises MemoryError if the traced memory increases by more than budget_bytes.

    Example:
        with MemoryBudget(1024 * 1024):   # 1 MiB budget
            small_operation()
        # Raises MemoryError if > 1 MiB was allocated net
    """

    def __init__(self, budget_bytes: int, nframe: int = 10) -> None:
        self.budget_bytes = budget_bytes
        self.nframe = nframe
        self._was_tracing = False
        self._start_snap: tracemalloc.Snapshot | None = None

    def __enter__(self) -> "MemoryBudget":
        self._was_tracing = tracemalloc.is_tracing()
        if not self._was_tracing:
            tracemalloc.start(self.nframe)
        gc.collect()
        self._start_snap = tracemalloc.take_snapshot()
        return self

    def __exit__(self, *args) -> None:
        gc.collect()
        end_snap = tracemalloc.take_snapshot()
        if not self._was_tracing:
            tracemalloc.stop()
        if self._start_snap is not None:
            diffs = end_snap.compare_to(self._start_snap, "lineno")
            net = sum(d.size_diff for d in diffs)
            if net > self.budget_bytes:
                top = sorted(diffs, key=lambda d: -d.size_diff)[:3]
                details = "; ".join(
                    f"{d.traceback[0].filename}:{d.traceback[0].lineno} +{d.size_diff/1024:.1f}KiB"
                    for d in top if d.size_diff > 0 and d.traceback
                )
                raise MemoryError(
                    f"Memory budget exceeded: {net/1024:.1f} KiB allocated "
                    f"(budget {self.budget_bytes/1024:.1f} KiB). Top: {details}"
                )


# ─────────────────────────────────────────────────────────────────────────────
# 5. Repeated-call leak monitor
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class LeakReport:
    iterations:   int
    total_growth: int    # bytes
    per_iter:     float  # bytes / iter
    top_sites:    list[AllocDiff] = field(default_factory=list)

    def __str__(self) -> str:
        return (f"iterations={self.iterations}  "
                f"growth={self.total_growth / 1024:.1f} KiB total  "
                f"{self.per_iter / 1024:.2f} KiB/iter")


def detect_leak(
    fn: Callable,
    iterations: int = 10,
    warmup: int = 2,
    nframe: int = 25,
) -> LeakReport:
    """
    Call fn repeatedly and check if memory grows with each iteration.
    Warmup runs are discarded before the baseline snapshot.

    Example:
        def maybe_leaks():
            global cache
            cache = {i: i*i for i in range(1000)}
        report = detect_leak(maybe_leaks, iterations=5)
        print(report)
        for site in report.top_sites:
            print(f"  {site}")
    """
    if not tracemalloc.is_tracing():
        tracemalloc.start(nframe)
        started = True
    else:
        started = False

    try:
        for _ in range(warmup):
            fn()
        gc.collect()
        snap_before = tracemalloc.take_snapshot()

        for _ in range(iterations):
            fn()
        gc.collect()
        snap_after = tracemalloc.take_snapshot()

    finally:
        if started:
            tracemalloc.stop()

    diffs = diff_snapshots(snap_before, snap_after, limit=10, only_growth=True)
    total = sum(d.size_diff for d in diffs)
    return LeakReport(
        iterations=iterations,
        total_growth=total,
        per_iter=total / iterations if iterations else 0,
        top_sites=diffs[:5],
    )


# ─────────────────────────────────────────────────────────────────────────────
# Demo
# ─────────────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=== tracemalloc demo ===")

    # ── basic trace + top ──────────────────────────────────────────────────────
    print("\n--- top allocations ---")
    tracemalloc.start(25)

    # Allocate some data
    big_list = [bytearray(1024) for _ in range(100)]
    nested = {str(i): list(range(100)) for i in range(50)}

    snap = take_snapshot(exclude_stdlib=True)
    tracemalloc.stop()

    print(f"  snapshot: {len(snap.statistics('lineno'))} allocation sites")
    for s in top_allocations(snap, limit=3):
        print(f"  {s}")

    # ── diff_snapshots ──────────────────────────────────────────────────────────
    print("\n--- diff_snapshots ---")
    tracemalloc.start(25)
    gc.collect()
    snap1 = take_snapshot()

    held = {i: bytearray(512) for i in range(200)}  # 200 * 512 bytes

    gc.collect()
    snap2 = take_snapshot()
    tracemalloc.stop()

    diffs = diff_snapshots(snap1, snap2, limit=5)
    for d in diffs:
        print(f"  {d}")

    # ── get_traced_memory ───────────────────────────────────────────────────────
    print("\n--- get_traced_memory ---")
    tracemalloc.start(10)
    stuff = [list(range(1000)) for _ in range(20)]
    cur, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    print(f"  current: {cur/1024:.1f} KiB  peak: {peak/1024:.1f} KiB")

    # ── MemoryBudget ────────────────────────────────────────────────────────────
    print("\n--- MemoryBudget ---")
    try:
        with MemoryBudget(10 * 1024 * 1024):   # 10 MiB — should pass
            tiny = list(range(100))
        print("  10 MiB budget: passed")
    except MemoryError as e:
        print(f"  unexpected: {e}")

    try:
        with MemoryBudget(1024):   # 1 KiB — should fail
            _big = bytearray(1024 * 1024)
        print("  1 KiB budget: unexpectedly passed")
    except MemoryError as e:
        print(f"  1 KiB budget exceeded: {str(e)[:80]}")

    # ── detect_leak ─────────────────────────────────────────────────────────────
    print("\n--- detect_leak ---")
    _accumulator: list = []
    def _leak_fn() -> None:
        _accumulator.extend(range(100))

    report = detect_leak(_leak_fn, iterations=5, warmup=1)
    print(f"  {report}")
    for site in report.top_sites[:2]:
        print(f"  {site}")

    print("\n=== done ===")

For the memray alternative — memray (PyPI) is a deterministic memory profiler that records allocations at the C level, captures native stack frames, exports flame graphs, and provides both CLI and programmatic APIs — use memray for deep production memory profiling with minimal code changes and rich visualization; use tracemalloc when you want zero third-party dependencies, need a pure Python context-manager-based approach, or are running in environments where installing native extensions is not possible. For the gc alternative — Python’s gc module exposes the garbage collector: gc.get_objects() lists all tracked objects; gc.get_referrers(obj) shows what holds a reference; gc.collect() forces a GC cycle — use gc to debug reference cycles and understand why specific objects are being retained; combine tracemalloc (to identify which allocation sites are growing) with gc.get_referrers() (to find what holds references to the leaked objects) for a complete leak investigation. The Claude Skills 360 bundle includes tracemalloc skill sets covering trace_memory() context manager, take_snapshot() with stdlib filtering, MemoryStat with top_allocations()/print_top(), AllocDiff with diff_snapshots(), MemoryBudget context manager with budget enforcement, and LeakReport with detect_leak() repeated-call leak detector. Start with the free tier to try memory leak detection patterns and tracemalloc pipeline code generation.