Don't leak partially-constructed modules when __init__ raises

[nstarman]

Module.__new__ adds the new instance to _currently_initialising, and _ModuleMeta.__call__ is responsible for removing it again. It did so via a tryself variable assigned from the result of super().__call__(...):

tryself = None
try:
    self = tryself = super().__call__(*args, **kwargs)
finally:
    if tryself is not None:
        _currently_initialising.remove(tryself)
    del tryself

If construction raises then the assignment is never reached, so tryself stays None and nothing is removed. As _IdSet deliberately holds a strong reference (to stop id being reallocated), every failed construction leaks its instance for the lifetime of the program.

class Boom(eqx.Module):
    x: int

    def __init__(self, x):
        self.x = x
        raise RuntimeError("boom")

for _ in range(100):
    try:
        Boom(1)
    except RuntimeError:
        pass

print(len(_currently_initialising))  # 100 before, 0 after

There are two routes in: a user __init__ raising, and BetterABCMeta.__call__ raising its __abstractvars__ check — the latter happening after the instance has been constructed.

Fix

We can't simply remove the offending instance on the error path, since type.__call__ doesn't hand back the partially-constructed object. Splitting __new__/__init__ explicitly in _ModuleMeta.__call__ would give us the instance, but it would bypass BetterABCMeta.__call__ and silently drop both of its abstract-class checks.

So instead: record the size of _currently_initialising beforehand, and unwind back to that mark if an exception propagates. _IdSet is backed by an insertion-ordered dict, and nested modules constructed during our __init__ remove their own entries, so anything left above the mark was added on our behalf.

The success path is unchanged in cost — a len and a remove, no allocation — which seemed worth preserving given the recent module-init perf work. (A weakref-based _IdSet would also fix this, but at the cost of a weakref allocation per instantiation.)

Tests

Three regression tests in tests/test_module.py, covering __init__ raising, a nested failure inside an outer __init__, and the AbstractVar route. Full suite passes locally (474 passed, 4 skipped).

🤖 Generated with Claude Code

[patrick-kidger]

Ah no! That's not good.

For the proposed fix though: will this be threadsafe?

[nstarman]

Good catch! No, the proposed fix wasn't thread safe.
Working on it....

[nstarman]

Claude tests are telling me that "losing slots plus the threading.local attribute lookup adds roughly 10% to Module() construction in a microbenchmark (~4360 → ~4806 ns for a 2-field module)". So this isn't optimal, but I haven't figured out a way to avoid this and solve the issue.