#!/usr/bin/env python3 """Find resource-handling bugs and scaffold correct context managers. `leaks` statically analyses Python source (it never imports or runs it) and reports resources opened outside `with` (files, sqlite3/psycopg connections, sockets, temp files/dirs, zip/tar archives, subprocess.Popen, thread/process pools, httpx/aiohttp/requests clients, urlopen), Lock/Semaphore .acquire() without a try/finally release, @contextmanager generators whose cleanup is skipped on exceptions or that yield more than once, except-blocks around `yield` that silently swallow errors, __exit__ methods that always return a truthy value (swallowing every exception), wrong __exit__ signatures, sync/async protocol mix-ups, `with threading.Lock():` on a fresh lock, and torch.no_grad()/suppress() objects created but never entered. `scaffold` prints a complete, typed, documented context manager (class, generator, async class, async generator, or a ready-made AI-engineering recipe) and verifies it with ast.parse before printing. Usage: python ctx_tool.py leaks app.py python ctx_tool.py leaks src/ --json cat app.py | python ctx_tool.py leaks - --strict # exit 1 on medium/high findings python ctx_tool.py leaks app.py --strict --fail-on high python ctx_tool.py rules # list rule ids python ctx_tool.py scaffold vector_tx --kind generator \\ --setup "conn = pool.getconn()" --teardown "pool.putconn(conn)" --yields conn python ctx_tool.py scaffold ModelSession --kind async-class \\ --setup "self.client = httpx.AsyncClient()" --teardown "await self.client.aclose()" \\ --import "import httpx" python ctx_tool.py scaffold --recipe atomic-write python ctx_tool.py scaffold --list-recipes Exit codes: 0 no findings (or findings without --strict), scaffold printed 1 --strict and at least one finding at or above --fail-on (default: medium) 2 usage error, unreadable file, Python syntax error, or generated code failed ast.parse Standard library only. Python 3.9+. """ import argparse import ast import json import os import re import sys import textwrap # --------------------------------------------------------------------------- # Rule catalogue # --------------------------------------------------------------------------- RULES = { "CTX001": ("resource-outside-with", "Resource opened outside `with` and not closed in a finally block"), "CTX002": ("acquire-without-finally", "Lock/Semaphore .acquire() without a try/finally release"), "CTX003": ("cm-yield-not-protected", "@contextmanager yield not inside try/finally: cleanup skipped on exception"), "CTX004": ("cm-yield-count", "@contextmanager generator must yield exactly once on every path"), "CTX005": ("exit-always-suppresses", "__exit__/__aexit__ returns a truthy value unconditionally (swallows all exceptions)"), "CTX006": ("enter-returns-none", "__enter__/__aenter__ returns None, so `with X() as x` binds None"), "CTX007": ("cm-except-swallows", "except around `yield` in @contextmanager does not re-raise (suppresses the body's exception)"), "CTX008": ("exit-signature", "__exit__/__aexit__ must accept (self, exc_type, exc, tb)"), "CTX009": ("sync-async-mismatch", "Sync/async protocol mix-up (async def __exit__, def __aenter__, @contextmanager on async def...)"), "CTX010": ("sync-with-async-resource", "Async-only resource used with plain `with` (needs `async with`)"), "CTX011": ("fresh-lock-in-with", "`with Lock():` creates a new lock each time and provides no mutual exclusion"), "CTX012": ("cm-never-entered", "Context manager object created as a bare statement and never entered (no effect)"), } SEV_RANK = {"low": 1, "medium": 2, "high": 3} # fully-qualified constructor -> (label, cleanup method, async-only) RESOURCES = { "open": ("file", "close", False), "io.open": ("file", "close", False), "codecs.open": ("file", "close", False), "gzip.open": ("gzip file", "close", False), "gzip.GzipFile": ("gzip file", "close", False), "bz2.open": ("bz2 file", "close", False), "lzma.open": ("lzma file", "close", False), "sqlite3.connect": ("sqlite3 connection", "close", False), "psycopg.connect": ("psycopg connection", "close", False), "psycopg2.connect": ("psycopg2 connection", "close", False), "socket.socket": ("socket", "close", False), "socket.create_connection": ("socket", "close", False), "socket.create_server": ("socket", "close", False), "tempfile.NamedTemporaryFile": ("temporary file", "close", False), "tempfile.TemporaryFile": ("temporary file", "close", False), "tempfile.SpooledTemporaryFile": ("temporary file", "close", False), "tempfile.TemporaryDirectory": ("temporary directory", "cleanup", False), "tempfile.mkdtemp": ("temporary directory (mkdtemp, not a context manager)", "shutil.rmtree", False), "zipfile.ZipFile": ("zip archive", "close", False), "tarfile.open": ("tar archive", "close", False), "tarfile.TarFile": ("tar archive", "close", False), "subprocess.Popen": ("subprocess", "wait", False), "concurrent.futures.ThreadPoolExecutor": ("thread pool", "shutdown", False), "concurrent.futures.ProcessPoolExecutor": ("process pool", "shutdown", False), "concurrent.futures.thread.ThreadPoolExecutor": ("thread pool", "shutdown", False), "multiprocessing.Pool": ("process pool", "close", False), "multiprocessing.pool.Pool": ("process pool", "close", False), "multiprocessing.pool.ThreadPool": ("thread pool", "close", False), "httpx.Client": ("httpx client", "close", False), "httpx.AsyncClient": ("httpx async client", "aclose", True), "aiohttp.ClientSession": ("aiohttp session", "close", True), "requests.Session": ("requests session", "close", False), "requests.session": ("requests session", "close", False), "requests.sessions.Session": ("requests session", "close", False), "urllib.request.urlopen": ("HTTP response", "close", False), } # Bare names recognised even when the import is not visible (snippets). BARE_FALLBACK = { "ThreadPoolExecutor": "concurrent.futures.ThreadPoolExecutor", "ProcessPoolExecutor": "concurrent.futures.ProcessPoolExecutor", "NamedTemporaryFile": "tempfile.NamedTemporaryFile", "TemporaryDirectory": "tempfile.TemporaryDirectory", "Popen": "subprocess.Popen", "AsyncClient": "httpx.AsyncClient", "ClientSession": "aiohttp.ClientSession", } CLOSERS = { "close", "aclose", "cleanup", "shutdown", "terminate", "kill", "wait", "communicate", "release", "rmtree", "remove", "unlink", "join", "__exit__", "__aexit__", "stop", "dispose", "disconnect", } OWNERSHIP_WRAPPERS = {"closing", "aclosing", "enter_context", "enter_async_context", "push", "push_async_exit"} LOCK_FACTORIES = { "threading.Lock", "threading.RLock", "threading.Semaphore", "threading.BoundedSemaphore", "threading.Condition", "asyncio.Lock", "asyncio.Semaphore", "asyncio.BoundedSemaphore", "asyncio.Condition", "multiprocessing.Lock", "multiprocessing.RLock", } NEVER_ENTERED = { "torch.no_grad", "torch.inference_mode", "torch.enable_grad", "torch.autocast", "torch.amp.autocast", "torch.cuda.amp.autocast", "torch.cpu.amp.autocast", "torch.autograd.profiler.profile", "torch.profiler.profile", "contextlib.suppress", "contextlib.redirect_stdout", "contextlib.redirect_stderr", "contextlib.nullcontext", "contextlib.chdir", "contextlib.ExitStack", "contextlib.AsyncExitStack", "warnings.catch_warnings", "decimal.localcontext", "numpy.errstate", "unittest.mock.patch", "unittest.mock.patch.object", "unittest.mock.patch.dict", } FUNC_TYPES = (ast.FunctionDef, ast.AsyncFunctionDef) SCOPE_TYPES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module, ast.Lambda) # --------------------------------------------------------------------------- # AST helpers # --------------------------------------------------------------------------- ALIASES = (("directory", "tmpdir"), ("client", "client"), ("session", "session"), ("connection", "conn"), ("socket", "sock"), ("pool", "pool"), ("subprocess", "proc"), ("archive", "archive"), ("response", "resp")) def alias_for(label, parent): """Variable name for the suggested `with ... as NAME` rewrite.""" if isinstance(parent, (ast.Assign, ast.AnnAssign)): tgt = parent.targets[0] if isinstance(parent, ast.Assign) else parent.target if isinstance(tgt, ast.Name): return tgt.id for key, alias in ALIASES: if key in label: return alias return "f" def unparse(node): try: return ast.unparse(node) except Exception: # pragma: no cover - defensive return "" def dotted(node): parts = [] while isinstance(node, ast.Attribute): parts.append(node.attr) node = node.value if isinstance(node, ast.Name): parts.append(node.id) return list(reversed(parts)) return None class Analyzer: def __init__(self, tree, source, filename): self.tree = tree self.lines = source.splitlines() self.filename = filename self.findings = [] self.parents = {} for node in ast.walk(tree): for child in ast.iter_child_nodes(node): self.parents[child] = node self.aliases = {} self.local_defs = set() for node in ast.walk(tree): if isinstance(node, ast.Import): for a in node.names: if a.asname: self.aliases[a.asname] = a.name else: top = a.name.split(".")[0] self.aliases[top] = top elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0: for a in node.names: if a.name != "*": self.aliases[a.asname or a.name] = node.module + "." + a.name elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): self.local_defs.add(node.name) # -- generic helpers --------------------------------------------------- def resolve(self, func): parts = dotted(func) if not parts: return None head = parts[0] if head in self.aliases: return ".".join([self.aliases[head]] + parts[1:]) if len(parts) == 1 and head in BARE_FALLBACK and head not in self.local_defs: return BARE_FALLBACK[head] return ".".join(parts) def add(self, rule, node, severity, message, fix=None): self.findings.append({ "file": self.filename, "line": getattr(node, "lineno", 0), "col": getattr(node, "col_offset", 0) + 1, "rule": rule, "name": RULES[rule][0], "severity": severity, "message": message, "fix": fix, }) def ancestors(self, node): while node in self.parents: node = self.parents[node] yield node def enclosing(self, node, types): for anc in self.ancestors(node): if isinstance(anc, types): return anc return None def stmt_of(self, node): cur = node while cur is not None and not isinstance(cur, ast.stmt): cur = self.parents.get(cur) return cur def walk_local(self, root): """Walk `root` without descending into nested functions/classes/lambdas.""" stack = list(ast.iter_child_nodes(root)) while stack: node = stack.pop() yield node if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)): continue stack.extend(ast.iter_child_nodes(node)) def in_field(self, node, owner, field): """True if `node` sits (transitively) inside owner..""" children = getattr(owner, field, None) or [] if not isinstance(children, list): children = [children] targets = set() for child in children: targets.add(child) cur = node while cur is not None and cur is not owner: if cur in targets: return True cur = self.parents.get(cur) return False # -- run --------------------------------------------------------------- def run(self): for node in ast.walk(self.tree): if isinstance(node, ast.Call): self.check_call(node) elif isinstance(node, FUNC_TYPES): self.check_function(node) elif isinstance(node, (ast.With, ast.AsyncWith)): self.check_with(node) self.findings.sort(key=lambda f: (f["line"], f["col"], f["rule"])) return self.findings # -- calls ------------------------------------------------------------- def resource_of(self, call): name = self.resolve(call.func) if name in RESOURCES: return name, RESOURCES[name] # Path(...).open() / some_path.open() if isinstance(call.func, ast.Attribute) and call.func.attr == "open": recv = call.func.value if isinstance(recv, ast.Call) and (self.resolve(recv.func) or "").split(".")[-1] in ("Path", "PurePath"): return "pathlib.Path.open", ("file", "close", False) parts = dotted(recv) if parts and re.search(r"path$", parts[-1], re.I): return "pathlib.Path.open", ("file", "close", False) return None, None def check_call(self, call): name = self.resolve(call.func) if isinstance(call.func, ast.Attribute) and call.func.attr == "acquire": self.check_acquire(call) return if name in NEVER_ENTERED: parent = self.parents.get(call) if isinstance(parent, ast.Expr): self.add("CTX012", call, "medium", f"`{unparse(call)}` builds a context manager and throws it away; it has no effect " "unless entered.", f"with {unparse(call)}:\n ... # code that should run under it") return rname, info = self.resource_of(call) if info: self.check_resource(call, rname, info) def check_resource(self, call, rname, info): label, closer, async_only = info node, parent = call, self.parents.get(call) if isinstance(parent, ast.Await): node, parent = parent, self.parents.get(parent) text = unparse(call) kw = "async with" if async_only else "with" alias = alias_for(label, parent) if isinstance(parent, ast.withitem): with_node = self.parents.get(parent) if async_only and isinstance(with_node, ast.With): self.add("CTX010", call, "high", f"{label} `{text}` only supports `async with`; plain `with` raises at runtime.", f"async with {text} as {alias}:\n ...") return if isinstance(parent, ast.keyword): parent = self.parents.get(parent) if isinstance(parent, ast.Call) and parent is not call: pname = (dotted(parent.func) or [""])[-1] if pname in OWNERSHIP_WRAPPERS: return if isinstance(self.parents.get(parent), ast.withitem): return # wrapper such as io.TextIOWrapper(open(...)) owned by the with if isinstance(parent, (ast.Return, ast.Yield, ast.YieldFrom, ast.Lambda)): return # ownership handed to the caller if rname == "tempfile.mkdtemp": fix = "with tempfile.TemporaryDirectory() as tmpdir:\n ... # removed automatically, even on error" else: fix = f"{kw} {text} as {alias}:\n ..." if isinstance(parent, (ast.Assign, ast.AnnAssign, ast.NamedExpr)): targets = parent.targets if isinstance(parent, ast.Assign) else [parent.target] for tgt in targets: self.check_owned(call, tgt, label, closer, fix, rname) return if isinstance(parent, ast.Expr): self.add("CTX001", call, "high", f"{label} `{text}` is created and discarded; it is never closed.", fix) elif isinstance(parent, ast.Attribute): self.add("CTX001", call, "medium", f"{label} `{text}` is used inline (`.{parent.attr}`) and left for the garbage collector " "to close (ResourceWarning; leaks on PyPy and in long-running workers).", fix) else: self.add("CTX001", call, "medium", f"{label} `{text}` is passed along without an owner that closes it.", fix) def check_owned(self, call, target, label, closer, fix, rname=""): names = [] if isinstance(target, (ast.Tuple, ast.List)): names = [unparse(e) for e in target.elts if isinstance(e, (ast.Name, ast.Attribute))] elif isinstance(target, (ast.Name, ast.Attribute)): names = [unparse(target)] if not names: self.add("CTX001", call, "low", f"{label} stored in `{unparse(target)}`; make sure some owner closes it.", fix) return is_self_attr = all(n.startswith("self.") for n in names) func = self.enclosing(call, FUNC_TYPES) if is_self_attr: scope = self.enclosing(call, ast.ClassDef) or func or self.tree else: scope = func or self.tree status = self.cleanup_status(scope, names, call) if status == "escapes" and rname == "tempfile.mkdtemp": status = "none" # passing a path string around does not transfer cleanup duty if status == "unsafe" and is_self_attr: status = "safe" # released by a close()/__exit__ method of the owning class shown = ", ".join(names) if status == "safe": return if status == "escapes": self.add("CTX001", call, "low", f"{label} `{shown}` is handed to another object/function; confirm that owner closes it.", fix) return if status == "unsafe": self.add("CTX001", call, "medium", f"{label} `{shown}` is closed, but not in a finally block: an exception before the " f"close leaks it.", fix) return if is_self_attr: self.add("CTX001", call, "medium", f"{label} `{shown}` is owned by the class but no close()/__exit__/__aexit__ releases it.", "def close(self) -> None:\n " + names[0] + "." + closer + "()\n\n" "def __enter__(self):\n return self\n\n" "def __exit__(self, exc_type, exc, tb) -> None:\n self.close()") elif scope is self.tree: self.add("CTX001", call, "low", f"module-level {label} `{shown}` is never closed; close it on shutdown " "(FastAPI lifespan, atexit) or create it inside a `with`.", fix) else: self.add("CTX001", call, "high", f"{label} `{shown}` is never closed (no {closer}() / with).", fix) def closes(self, call, names): """Does this Call release one of `names`?""" func = call.func if isinstance(func, ast.Attribute): if func.attr in CLOSERS and unparse(func.value) in names: return True if func.attr in CLOSERS and any(unparse(a) in names for a in call.args): return True if func.attr in ("enter_context", "enter_async_context", "callback", "push_async_callback", "finalize"): for a in call.args: t = unparse(a) if t in names or any(t.startswith(n + ".") for n in names): return True elif isinstance(func, ast.Name): if func.id in CLOSERS and any(unparse(a) in names for a in call.args): return True if func.id in ("closing", "aclosing", "finalize"): return any(unparse(a) in names or any(unparse(a).startswith(n + ".") for n in names) for a in call.args) return False def cleanup_status(self, scope, names, origin): status = "none" for node in ast.walk(scope): if isinstance(node, ast.Try) and node.finalbody: for stmt in node.finalbody: for sub in ast.walk(stmt): if isinstance(sub, ast.Call) and self.closes(sub, names): return "safe" elif isinstance(node, (ast.With, ast.AsyncWith)): for item in node.items: ce = item.context_expr if unparse(ce) in names: return "safe" if isinstance(ce, ast.Call) and self.closes(ce, names): return "safe" if isinstance(ce, ast.Call) and any(unparse(a) in names for a in ce.args): return "safe" elif isinstance(node, ast.Return) and node.value is not None: vals = node.value.elts if isinstance(node.value, ast.Tuple) else [node.value] if any(unparse(v) in names for v in vals): return "safe" elif isinstance(node, ast.Assign) and unparse(node.value) in names: if any(unparse(t).startswith("self.") or isinstance(t, (ast.Attribute, ast.Subscript)) for t in node.targets): return "safe" # ownership moved to an object attribute/container elif isinstance(node, ast.Call) and node is not origin: if self.closes(node, names): if self.closing_in_handler(node): status = "unsafe" if status != "escapes" else status elif status == "none": status = "unsafe" elif status == "none" and any(unparse(a) in names for a in node.args) \ and not (isinstance(node.func, ast.Attribute) and unparse(node.func.value) in names): status = "escapes" return status def closing_in_handler(self, node): return any(isinstance(a, ast.ExceptHandler) for a in self.ancestors(node)) def check_acquire(self, call): recv = unparse(call.func.value) if isinstance(self.parents.get(call), ast.withitem): return node = call if isinstance(self.parents.get(node), ast.Await): node = self.parents[node] stmt = self.stmt_of(node) if stmt is None: return scope = self.enclosing(call, FUNC_TYPES) or self.tree if getattr(scope, "name", "") in ("__enter__", "__aenter__"): cls = self.parents.get(scope) if isinstance(cls, ast.ClassDef) and self.class_releases(cls, recv): return # acquired on enter, released on exit: the protocol is the try/finally names = {recv} if isinstance(stmt, (ast.Assign, ast.AnnAssign)) and self.parents.get(node) is stmt: tgts = stmt.targets if isinstance(stmt, ast.Assign) else [stmt.target] names |= {unparse(t) for t in tgts} released_in_finally, released_anywhere, inside_try = False, False, False for t in ast.walk(scope): if isinstance(t, ast.Try) and t.finalbody: for s in t.finalbody: for sub in ast.walk(s): if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Attribute) \ and sub.func.attr == "release" \ and (unparse(sub.func.value) in names or any(unparse(a) in names for a in sub.args)): released_in_finally = True if self.in_field(stmt, t, "body"): inside_try = True if isinstance(t, ast.Call) and isinstance(t.func, ast.Attribute) and t.func.attr == "release" \ and unparse(t.func.value) in names: released_anywhere = True fix = (f"async with {recv}:\n ..." if isinstance(node, ast.Await) else f"with {recv}:\n ...") if released_in_finally and inside_try: self.add("CTX002", call, "low", f"`{recv}.acquire()` is inside the try whose finally releases it: if acquire() raises " "(timeout, cancellation) the finally releases a lock you never held. Acquire before `try:`.", fix) elif released_in_finally: return elif released_anywhere: self.add("CTX002", call, "high", f"`{recv}` is released, but not in a finally block: an exception between acquire and " "release deadlocks every later caller.", fix) else: self.add("CTX002", call, "high", f"`{recv}.acquire()` has no matching release() in a finally block.", fix) def class_releases(self, cls, recv): for item in cls.body: if isinstance(item, FUNC_TYPES) and item.name in ("__exit__", "__aexit__"): for sub in ast.walk(item): if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Attribute) \ and sub.func.attr == "release" and unparse(sub.func.value) == recv: return True return False # -- with statements --------------------------------------------------- def check_with(self, node): for item in node.items: ce = item.context_expr if isinstance(ce, ast.Call): name = self.resolve(ce.func) if name in LOCK_FACTORIES: kw = "async with" if name.startswith("asyncio.") else "with" self.add("CTX011", ce, "high", f"`{unparse(ce)}` creates a brand-new lock on every entry, so nothing is ever " "excluded. Create the lock once (module, __init__, or app state) and reuse it.", f"_lock = {unparse(ce)} # once, shared\n...\n{kw} _lock:\n ...") # -- functions --------------------------------------------------------- def decorator_names(self, fn): out = [] for d in fn.decorator_list: target = d.func if isinstance(d, ast.Call) else d parts = dotted(target) if parts: out.append(parts[-1]) return out def check_function(self, fn): decos = self.decorator_names(fn) is_async = isinstance(fn, ast.AsyncFunctionDef) if "contextmanager" in decos or "asynccontextmanager" in decos: if "contextmanager" in decos and is_async: self.add("CTX009", fn, "high", f"`{fn.name}` is `async def` but decorated with @contextmanager; use " "@asynccontextmanager and `async with`.", "@asynccontextmanager\nasync def " + fn.name + "(...):") if "asynccontextmanager" in decos and not is_async: self.add("CTX009", fn, "high", f"`{fn.name}` is a plain def decorated with @asynccontextmanager; make it `async def` " "or use @contextmanager.", "@contextmanager\ndef " + fn.name + "(...):") self.check_cm_generator(fn) cls = self.parents.get(fn) if isinstance(cls, ast.ClassDef): self.check_dunder(fn, cls, is_async) def check_cm_generator(self, fn): yields = sorted((n for n in self.walk_local(fn) if isinstance(n, (ast.Yield, ast.YieldFrom))), key=lambda n: (n.lineno, n.col_offset)) if not yields: self.add("CTX004", fn, "high", f"@contextmanager `{fn.name}` never yields: entering it raises " "\"generator didn't yield\" (or TypeError).", "try:\n yield resource\nfinally:\n cleanup()") return for y in yields: loop = None for anc in self.ancestors(y): if anc is fn: break if isinstance(anc, (ast.For, ast.AsyncFor, ast.While)): loop = anc break if loop is not None: self.add("CTX004", y, "high", f"`yield` inside a loop in @contextmanager `{fn.name}`: a second iteration raises " "RuntimeError(\"generator didn't stop\").", "Build the value first, then yield exactly once outside the loop.") reported = False for i, a in enumerate(yields): for b in yields[i + 1:]: if not self.exclusive(a, b, fn) and not reported: reported = True self.add("CTX004", b, "high", f"@contextmanager `{fn.name}` can yield twice (lines {a.lineno} and {b.lineno}); " "a yield in an except/finally after the main yield raises " "RuntimeError(\"generator didn't stop after throw()\").", "try:\n yield value\nexcept SomeError:\n handle()\n raise\nfinally:\n cleanup()") for y in yields: self.check_yield_protection(fn, y) def exclusive(self, a, b, fn): pa = [a] + [x for x in self.ancestors(a)] pb = [b] + [x for x in self.ancestors(b)] set_b = {id(x): x for x in pb} lca, child_a = None, None prev = a for x in pa[1:]: if id(x) in set_b: lca = x child_a = prev break prev = x if lca is None: return False child_b = b for x in pb[1:]: if x is lca: break child_b = x if isinstance(lca, ast.If): return (child_a in lca.body) != (child_b in lca.body) if isinstance(lca, ast.IfExp): return child_a is not child_b and lca.test not in (child_a, child_b) match_case = getattr(ast, "match_case", None) if match_case is not None and isinstance(child_a, match_case) and isinstance(child_b, match_case): return child_a is not child_b return False def check_yield_protection(self, fn, y): stmt = self.stmt_of(y) protected = False for anc in self.ancestors(y): if anc is fn: break if isinstance(anc, ast.Try) and self.in_field(y, anc, "body"): if anc.finalbody: protected = True if anc.handlers: protected = True self.check_handlers(fn, anc, y) elif isinstance(anc, (ast.With, ast.AsyncWith)) and self.in_field(y, anc, "body"): protected = True if protected: return # Is there any code that runs after the yield (i.e., cleanup)? after = [] cur = stmt while cur is not None and cur is not fn: parent = self.parents.get(cur) for field in ("body", "orelse", "finalbody"): block = getattr(parent, field, None) if isinstance(block, list) and cur in block: after.extend(block[block.index(cur) + 1:]) cur = parent after = [s for s in after if not (isinstance(s, ast.Return) and s.value is None) and not isinstance(s, (ast.Pass, ast.Raise))] if not after: return cleanup = "\n".join(" " + unparse(s).splitlines()[0] for s in after[:3]) yv = unparse(y.value) if getattr(y, "value", None) is not None else "" self.add("CTX003", y, "high", f"cleanup after `yield` in `{fn.name}` (line {after[0].lineno}) is skipped when the with-body " "raises: the exception is thrown in at the yield and never reaches it.", f"try:\n yield {yv}".rstrip() + f"\nfinally:\n{cleanup}") def check_handlers(self, fn, try_node, y): for h in try_node.handlers: reraises = any(isinstance(n, ast.Raise) for n in self.walk_local(h)) if reraises: continue caught = unparse(h.type) if h.type is not None else "everything" self.add("CTX007", h, "medium", f"`except {caught}` around the yield in `{fn.name}` does not re-raise, so exceptions " "raised inside the with-block are silently suppressed. Intentional only for " "suppress()-style managers.", f"except {unparse(h.type) if h.type is not None else 'BaseException'}:\n rollback()\n raise") def check_dunder(self, fn, cls, is_async): name = fn.name if name in ("__enter__", "__exit__") and is_async: self.add("CTX009", fn, "high", f"`async def {name}` in `{cls.name}`: the sync protocol never awaits it. Use " f"`{name.replace('__', '__a', 1)}` and `async with`.", f"async def {name.replace('__', '__a', 1)}(self, ...):") if name in ("__aenter__", "__aexit__") and not is_async: self.add("CTX009", fn, "high", f"`{name}` in `{cls.name}` must be `async def`; `async with` awaits its result.", f"async def {name}(self, ...):") if name in ("__exit__", "__aexit__"): npos = len(getattr(fn.args, "posonlyargs", [])) + len(fn.args.args) if fn.args.vararg is None and npos != 4: self.add("CTX008", fn, "high", f"`{cls.name}.{name}` takes {npos} positional parameters; Python calls it with " "(self, exc_type, exc, tb), so leaving the with-block raises TypeError.", f"def {name}(self, exc_type, exc, tb) -> None:") for r in self.walk_local(fn): if isinstance(r, ast.Return) and r.value is not None and self.truthy(r.value) \ and self.unconditional(r, fn): self.add("CTX005", r, "high", f"`{cls.name}.{name}` always returns `{unparse(r.value)}`, which suppresses every " "exception from the with-body (including real bugs such as KeyError, and " "KeyboardInterrupt). Return None/False, or True only for specific handled types.", "def __exit__(self, exc_type, exc, tb) -> bool:\n self.close()\n" " return exc_type is not None and issubclass(exc_type, ExpectedError)") if name in ("__enter__", "__aenter__"): returns_value = any(isinstance(r, ast.Return) and r.value is not None for r in self.walk_local(fn)) has_yield = any(isinstance(r, (ast.Yield, ast.YieldFrom)) for r in self.walk_local(fn)) if not returns_value and not has_yield: used_as = self.bound_with_as(cls.name) sev = "medium" if used_as else "low" where = f" (line {used_as} binds it with `as`)" if used_as else "" self.add("CTX006", fn, sev, f"`{cls.name}.{name}` returns None, so `with {cls.name}(...) as x` binds x to None{where}.", f"{'async ' if is_async else ''}def {name}(self) -> \"{cls.name}\":\n ...\n return self") def bound_with_as(self, cls_name): for node in ast.walk(self.tree): if isinstance(node, (ast.With, ast.AsyncWith)): for item in node.items: ce = item.context_expr if item.optional_vars is not None and isinstance(ce, ast.Call) \ and (dotted(ce.func) or [""])[-1] == cls_name: return node.lineno return 0 @staticmethod def truthy(value): if isinstance(value, ast.Constant): return bool(value.value) return isinstance(value, ast.Name) and value.id == "self" def unconditional(self, node, fn): for anc in self.ancestors(node): if anc is fn: return True if isinstance(anc, (ast.If, ast.IfExp, ast.For, ast.AsyncFor, ast.While, ast.ExceptHandler)): return False if isinstance(anc, ast.Try) and not self.in_field(node, anc, "finalbody") and anc.handlers: return False if type(anc).__name__ in ("Match", "match_case"): return False return True # --------------------------------------------------------------------------- # leaks command # --------------------------------------------------------------------------- def iter_files(paths): for p in paths: if p == "-" or os.path.isfile(p): yield p, True elif os.path.isdir(p): for root, dirs, files in os.walk(p): dirs[:] = sorted(d for d in dirs if not d.startswith(".") and d not in ("__pycache__", "node_modules", "venv", ".venv")) for f in sorted(files): if f.endswith(".py"): yield os.path.join(root, f), False else: raise FileNotFoundError(p) def analyse_source(source, filename): tree = ast.parse(source, filename=filename) return Analyzer(tree, source, filename).run() def cmd_leaks(args): findings, warnings, scanned = [], [], 0 try: files = list(iter_files(args.paths)) except FileNotFoundError as e: print(f"error: no such file or directory: {e}", file=sys.stderr) return 2 for path, explicit in files: try: if path == "-": source = sys.stdin.read() else: with open(path, encoding="utf-8") as fh: source = fh.read() except (OSError, UnicodeDecodeError) as e: if explicit: print(f"error: cannot read {path}: {e}", file=sys.stderr) return 2 warnings.append(f"skipped {path}: {e}") continue name = "" if path == "-" else path try: findings.extend(analyse_source(source, name)) except SyntaxError as e: msg = f"{name}:{e.lineno}: syntax error: {e.msg}" if explicit: print(f"error: {msg}", file=sys.stderr) return 2 warnings.append("skipped " + msg) continue scanned += 1 if args.min_severity: findings = [f for f in findings if SEV_RANK[f["severity"]] >= SEV_RANK[args.min_severity]] counts = {s: sum(1 for f in findings if f["severity"] == s) for s in ("high", "medium", "low")} failing = [f for f in findings if SEV_RANK[f["severity"]] >= SEV_RANK[args.fail_on]] if args.json: print(json.dumps({"files_scanned": scanned, "findings": findings, "summary": counts, "warnings": warnings}, indent=2)) else: for w in warnings: print(f"warning: {w}", file=sys.stderr) for f in findings: print(f"{f['file']}:{f['line']}:{f['col']} {f['severity'].upper():<6} {f['rule']} {f['name']}") for line in textwrap.wrap(f["message"], 96): print(f" {line}") if f["fix"]: print(" suggested fix:") print(textwrap.indent(f["fix"], " ")) print() if findings: print(f"{len(findings)} finding(s) in {scanned} file(s): " f"{counts['high']} high, {counts['medium']} medium, {counts['low']} low") else: print(f"No resource-handling issues found in {scanned} file(s).") return 1 if (args.strict and failing) else 0 def cmd_rules(args): if args.json: print(json.dumps({rid: {"name": n, "summary": s} for rid, (n, s) in RULES.items()}, indent=2)) else: for rid, (n, s) in RULES.items(): print(f"{rid} {n:<26} {s}") return 0 # --------------------------------------------------------------------------- # scaffold command # --------------------------------------------------------------------------- def camel(name): if "_" not in name and name[:1].isupper(): return name return "".join(p[:1].upper() + p[1:] for p in name.split("_") if p) def snake(name): return re.sub(r"(? None: self._active = False {adef} __{a}enter__(self) -> {ret_type}: if self._active: raise RuntimeError("{cls} is not re-entrant; create a new instance") {block(setup, 8)} self._active = True return {ret} {adef} __{a}exit__( self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], tb: Optional[TracebackType], ) -> None: try: {block(teardown, 12)} finally: self._active = False ''' return "\n".join(header) + body is_async = kind == "async-generator" deco = "asynccontextmanager" if is_async else "contextmanager" it = "AsyncIterator" if is_async else "Iterator" adef = "async def" if is_async else "def" aw = "async with" if is_async else "with" fn = snake(name) if name[:1].isupper() else name ytype = "Any" if yields else "None" header += [f"from contextlib import {deco}", f"from typing import Any, {it}"] + list(imports) + ["", "", ""] body = f'''@{deco} {adef} {fn}() -> {it}[{ytype}]: """Run setup, yield {yields or 'None'}, and always run teardown. Setup: {setup_desc} Teardown: {teardown_desc} Setup runs before ``try`` so a failed setup does not trigger teardown of something that was never acquired. The body's exception is re-raised. Example:: {aw} {fn}(){' as value' if yields else ''}: ... """ {block(setup, 4)} try: yield {yields or 'None'} finally: {block(teardown, 8)} ''' return "\n".join(header) + body RECIPES = {} def recipe(name, default): def deco(fn): RECIPES[name] = (default, fn.__doc__.strip(), fn) return fn return deco @recipe("timer", "Timer") def _r_timer(): """Wall-clock timer (sync class, also usable as a decorator) for LLM calls, retrieval, batches.""" return '''"""Timer context manager / decorator for latency measurements.""" from __future__ import annotations import logging import time from contextlib import ContextDecorator from types import TracebackType from typing import Optional, Type logger = logging.getLogger(__name__) class __NAME__(ContextDecorator): """Measure elapsed wall time with time.perf_counter(). ``elapsed`` is set even when the body raises, and the exception propagates. Example:: with __NAME__("retrieval") as t: docs = retriever.search(query, k=8) print(f"{t.label}: {t.elapsed * 1000:.1f} ms") @__NAME__("embed_batch") def embed_batch(texts: list[str]) -> list[list[float]]: ... """ def __init__(self, label: str = "block", log: bool = True) -> None: self.label = label self.log = log self.start: Optional[float] = None self.elapsed: float = 0.0 def __enter__(self) -> "__NAME__": self.start = time.perf_counter() return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], tb: Optional[TracebackType], ) -> None: assert self.start is not None self.elapsed = time.perf_counter() - self.start if self.log: status = "error" if exc_type else "ok" logger.info("%s took %.1f ms (%s)", self.label, self.elapsed * 1000, status) ''' @recipe("temp-env", "temp_env") def _r_env(): """Temporarily set/unset environment variables (API keys, model endpoints) and restore them.""" return '''"""Temporarily override environment variables.""" from __future__ import annotations import os from contextlib import contextmanager from typing import Iterator, Optional @contextmanager def __NAME__(**overrides: Optional[str]) -> Iterator[None]: """Set env vars for the duration of the block; ``None`` unsets a variable. The previous values (including "was not set") are restored in ``finally``. os.environ is process-global: do not use this from concurrent threads/tasks. Example:: with __NAME__(OPENAI_API_KEY="sk-test", HTTPS_PROXY=None): client = make_client() """ saved = {key: os.environ.get(key) for key in overrides} try: for key, value in overrides.items(): if value is None: os.environ.pop(key, None) else: os.environ[key] = value yield finally: for key, old in saved.items(): if old is None: os.environ.pop(key, None) else: os.environ[key] = old ''' @recipe("cwd", "working_directory") def _r_cwd(): """Temporarily change the working directory (contextlib.chdir on 3.11+), e.g. for agent sandboxes.""" return '''"""Temporarily change the current working directory.""" from __future__ import annotations import os from contextlib import contextmanager from pathlib import Path from typing import Iterator, Union @contextmanager def __NAME__(path: Union[str, "os.PathLike[str]"]) -> Iterator[Path]: """chdir into ``path`` and always chdir back. Python 3.11+ ships ``contextlib.chdir`` with the same behaviour. The cwd is process-global, so do not use this from concurrent threads or asyncio tasks; pass ``cwd=`` to subprocess.run instead where you can. Example:: with __NAME__(sandbox_dir) as here: subprocess.run(["python", "tool.py"], check=True) """ previous = os.getcwd() os.chdir(path) try: yield Path(os.getcwd()) finally: os.chdir(previous) ''' @recipe("gpu-inference", "inference_session") def _r_gpu(): """eval() + torch.inference_mode() + optional autocast, restoring train mode and freeing CUDA cache.""" return '''"""Inference context for PyTorch models.""" from __future__ import annotations import gc from contextlib import ExitStack, contextmanager from typing import Iterator, Optional import torch @contextmanager def __NAME__( model: torch.nn.Module, *, amp_dtype: Optional[torch.dtype] = None, empty_cache: bool = False, ) -> Iterator[torch.nn.Module]: """Put ``model`` in eval mode under inference_mode (and autocast), then restore. - The original ``model.training`` flag is restored even on error, so a validation pass inside a training loop does not leave dropout disabled. - ``amp_dtype`` (torch.float16 / torch.bfloat16) enables autocast on the model's device type. - ``empty_cache`` runs gc + torch.cuda.empty_cache() on exit; this returns cached blocks to the driver but does not free tensors you still reference. Example:: with __NAME__(model, amp_dtype=torch.bfloat16) as m: logits = m(input_ids) """ was_training = model.training first_param = next(model.parameters(), None) device_type = first_param.device.type if first_param is not None else "cpu" with ExitStack() as stack: model.eval() stack.callback(model.train, was_training) if empty_cache: stack.callback(_free_cuda_cache) stack.enter_context(torch.inference_mode()) if amp_dtype is not None: stack.enter_context(torch.autocast(device_type=device_type, dtype=amp_dtype)) yield model def _free_cuda_cache() -> None: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() ''' @recipe("atomic-write", "atomic_write") def _r_atomic(): """Write a file atomically (temp file in same dir + fsync + os.replace) for checkpoints and indexes.""" return '''"""Atomic file writes: readers never see a half-written file.""" from __future__ import annotations import os import tempfile from contextlib import contextmanager from pathlib import Path from typing import IO, Any, Iterator, Union @contextmanager def __NAME__( path: Union[str, "os.PathLike[str]"], mode: str = "w", encoding: Union[str, None] = "utf-8", ) -> Iterator[IO[Any]]: """Yield a temp file next to ``path``; on success fsync and os.replace it into place. On any exception the temp file is deleted and ``path`` is untouched, so a crash mid-save never corrupts the previous checkpoint/index/manifest. Use ``mode="wb"`` (encoding is then ignored) for torch.save / numpy.save. Example:: with __NAME__("ckpt/model.pt", mode="wb") as f: torch.save(model.state_dict(), f) """ target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) binary = "b" in mode fd, tmp_name = tempfile.mkstemp(dir=target.parent, prefix=f".{target.name}.", suffix=".tmp") try: with os.fdopen(fd, mode, encoding=None if binary else encoding) as fh: yield fh fh.flush() os.fsync(fh.fileno()) os.replace(tmp_name, target) except BaseException: try: os.unlink(tmp_name) except FileNotFoundError: pass raise ''' @recipe("semaphore", "RateLimiter") def _r_sem(): """Async concurrency + rate limiter (asyncio.Semaphore with min spacing) for LLM/embedding fan-out.""" return '''"""Async concurrency limiter for API fan-out.""" from __future__ import annotations import asyncio import time from types import TracebackType from typing import Optional, Type class __NAME__: """Cap in-flight requests and enforce a minimum interval between starts. Create ONE instance and share it across tasks; ``async with`` it around each call. The slot is released even if the call raises or the task is cancelled. Example:: limiter = __NAME__(max_concurrency=8, min_interval=0.05) async def embed(text: str) -> list[float]: async with limiter: resp = await client.embeddings.create(model="text-embedding-3-small", input=text) return resp.data[0].embedding vectors = await asyncio.gather(*(embed(t) for t in texts)) """ def __init__(self, max_concurrency: int = 8, min_interval: float = 0.0) -> None: if max_concurrency < 1: raise ValueError("max_concurrency must be >= 1") self._sem = asyncio.Semaphore(max_concurrency) self._min_interval = min_interval self._last_start = 0.0 self._start_lock = asyncio.Lock() async def __aenter__(self) -> "__NAME__": await self._sem.acquire() try: if self._min_interval: async with self._start_lock: wait = self._last_start + self._min_interval - time.monotonic() if wait > 0: await asyncio.sleep(wait) self._last_start = time.monotonic() except BaseException: self._sem.release() raise return self async def __aexit__( self, exc_type: Optional[Type[BaseException]], exc: Optional[BaseException], tb: Optional[TracebackType], ) -> None: self._sem.release() ''' @recipe("span", "span") def _r_span(): """OpenTelemetry-style tracing span with nesting via contextvars, attributes, error status, duration.""" return '''"""Minimal tracing spans (OpenTelemetry-style) for LLM pipelines.""" from __future__ import annotations import contextvars import time import uuid from contextlib import contextmanager from dataclasses import dataclass, field from typing import Any, Callable, Dict, Iterator, List, Optional @dataclass class Span: name: str trace_id: str span_id: str parent_id: Optional[str] attributes: Dict[str, Any] = field(default_factory=dict) start: float = 0.0 duration_ms: float = 0.0 status: str = "ok" error: Optional[str] = None def set(self, key: str, value: Any) -> None: self.attributes[key] = value _current: contextvars.ContextVar[Optional[Span]] = contextvars.ContextVar("current_span", default=None) FINISHED: List[Span] = [] @contextmanager def __NAME__( name: str, exporter: Callable[[Span], None] = FINISHED.append, **attributes: Any, ) -> Iterator[Span]: """Open a child of the current span; record duration, attributes and errors. contextvars makes nesting correct across threads and asyncio tasks. The exception is recorded and re-raised, never swallowed. Example:: with __NAME__("rag.answer", query=q) as s: with __NAME__("retrieve", k=8): docs = retriever.search(q, k=8) with __NAME__("llm.generate", model="gpt-4o-mini") as g: resp = llm(prompt) g.set("tokens.total", resp.usage.total_tokens) """ parent = _current.get() s = Span( name=name, trace_id=parent.trace_id if parent else uuid.uuid4().hex, span_id=uuid.uuid4().hex[:16], parent_id=parent.span_id if parent else None, attributes=dict(attributes), start=time.perf_counter(), ) token = _current.set(s) try: yield s except BaseException as exc: s.status = "error" s.error = f"{type(exc).__name__}: {exc}" raise finally: s.duration_ms = (time.perf_counter() - s.start) * 1000 _current.reset(token) exporter(s) ''' @recipe("seed", "seeded") def _r_seed(): """Deterministic seeds for random/numpy/torch inside a block, restoring the previous RNG state.""" return '''"""Deterministic RNG block for evals and tests.""" from __future__ import annotations import random from contextlib import contextmanager from typing import Any, Dict, Iterator @contextmanager def __NAME__(seed: int = 0) -> Iterator[int]: """Seed Python, NumPy and PyTorch (if installed); restore their states on exit. Restoring keeps the rest of the program's randomness unaffected, so a seeded eval block inside a training loop does not change later shuffles. Note: full GPU determinism also needs torch.use_deterministic_algorithms(True) and CUBLAS_WORKSPACE_CONFIG; that is process-wide and not toggled here. Example:: with __NAME__(1234): sample = random.sample(dataset, 100) """ saved: Dict[str, Any] = {"random": random.getstate()} try: import numpy as np except ImportError: np = None try: import torch except ImportError: torch = None if np is not None: saved["numpy"] = np.random.get_state() if torch is not None: saved["torch"] = torch.random.get_rng_state() if torch.cuda.is_available(): saved["cuda"] = torch.cuda.get_rng_state_all() try: random.seed(seed) if np is not None: np.random.seed(seed) if torch is not None: torch.manual_seed(seed) # also seeds all CUDA devices yield seed finally: random.setstate(saved["random"]) if np is not None: np.random.set_state(saved["numpy"]) if torch is not None: torch.random.set_rng_state(saved["torch"]) if "cuda" in saved: torch.cuda.set_rng_state_all(saved["cuda"]) ''' def cmd_scaffold(args): if args.list_recipes: if args.json: print(json.dumps({k: {"default_name": v[0], "summary": v[1]} for k, v in RECIPES.items()}, indent=2)) else: for k, (default, doc, _) in RECIPES.items(): print(f"{k:<14} {default:<20} {doc}") return 0 if args.recipe: default, _, fn = RECIPES[args.recipe] name = args.name or default if not name.isidentifier(): print(f"error: NAME {name!r} is not a valid Python identifier", file=sys.stderr) return 2 if default[:1].isupper(): name = camel(name) code = fn().replace("__NAME__", name) kind = "recipe:" + args.recipe else: if not args.name: print("error: scaffold needs NAME (or --recipe / --list-recipes)", file=sys.stderr) return 2 if not args.name.isidentifier(): print(f"error: NAME {args.name!r} is not a valid Python identifier", file=sys.stderr) return 2 code = scaffold_kind(args.name, args.kind, args.setup, args.teardown, args.yields, args.imports or []) kind = args.kind try: ast.parse(code) except SyntaxError as e: print(f"error: generated code does not parse (line {e.lineno}: {e.msg}); check --setup/--teardown/" f"--yields/--import snippets.\n\n{code}", file=sys.stderr) return 2 if args.json: print(json.dumps({"kind": kind, "code": code, "verified": "ast.parse"}, indent=2)) else: print(code, end="" if code.endswith("\n") else "\n") return 0 # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def build_parser(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) sub = ap.add_subparsers(dest="command") lk = sub.add_parser("leaks", help="find resource leaks and context-manager bugs (static, never runs code)") lk.add_argument("paths", nargs="+", help="Python files, directories (recursed), or - for stdin") lk.add_argument("--json", action="store_true", help="machine-readable output") lk.add_argument("--strict", action="store_true", help="exit 1 if any finding is at or above --fail-on") lk.add_argument("--fail-on", choices=["low", "medium", "high"], default="medium", help="severity that fails --strict (default: medium)") lk.add_argument("--min-severity", choices=["low", "medium", "high"], help="hide findings below this severity") ru = sub.add_parser("rules", help="list rule ids") ru.add_argument("--json", action="store_true") sc = sub.add_parser("scaffold", help="print a typed, documented context manager (verified with ast.parse)") sc.add_argument("name", nargs="?", help="function or class name (CamelCase is applied for class kinds)") sc.add_argument("--kind", choices=["class", "generator", "async-class", "async-generator"], default="generator") sc.add_argument("--setup", default="", help="code run on entry (use \\n for several lines)") sc.add_argument("--teardown", default="", help="code always run on exit (use \\n for several lines)") sc.add_argument("--yields", default="", help="expression yielded / returned from __enter__ (default: None / self)") sc.add_argument("--import", dest="imports", action="append", help="import line to add (repeatable)") sc.add_argument("--recipe", choices=sorted(RECIPES), help="emit a ready-made recipe instead") sc.add_argument("--list-recipes", action="store_true", help="list recipes") sc.add_argument("--json", action="store_true", help="wrap output in JSON") return ap def main(argv=None): ap = build_parser() args = ap.parse_args(argv) if args.command == "leaks": return cmd_leaks(args) if args.command == "rules": return cmd_rules(args) if args.command == "scaffold": return cmd_scaffold(args) ap.print_help(sys.stderr) return 2 if __name__ == "__main__": sys.exit(main())