Stacklevel for ResourceWarning source

I’m working on adding long-overdue ResourceWarnings to my package, and discovered the wonderful source argument added in bpo-26567.

However, I’m running into the always-tricky stacklevel issue with wranings because my objects are allocated by an internal method, whereas the useful stack frame is the user’s own code one level up from that.

Is there an equivalent to the stacklevel argument for the warning itself that applies to the tracemalloc frame for source?

Or perhaps I could rewrite my method to skip the frame? I notice that if I replace the method with a functools.partial I get what I want, but I can’t quite do that in my case.

Here’s an example:

import warnings
from functools import partial


class Branch:
    def __init__(self, root):
        self.root = root
        self._closed = False

    def close(self):
        self._closed = True

    def __del__(self):
        if not self._closed:
            warnings.warn(
                f"unclosed branch {self}",
                ResourceWarning,
                source=self,
            )


class Root:
    def __init__(self):
        self._branches = []
        self.partial_branch = partial(Branch, self)

    def branch(self):
        b = Branch(self)  # this isn't a useful line to warn about
        self._branches.append(b)
        return b


def main():
    r = Root()
    # want a warning for both of these lines
    method_allocated = r.branch()  # warning about this branch is above in Root.branch
    partial_allocated = r.partial_branch()  # good warning


if __name__ == "__main__":
    import tracemalloc

    tracemalloc.start()
    warnings.simplefilter("default")
    main()

whilch gives the output:

test.py:19: ResourceWarning: unclosed branch <__main__.Branch object at 0x10b68d3d0>
  warnings.warn(
Object allocated at (most recent call last):
  File "test.py", lineno 42
    partial_allocated = r.partial_branch()  # good warning
test.py:19: ResourceWarning: unclosed branch <__main__.Branch object at 0x10b68d400>
  warnings.warn(
Object allocated at (most recent call last):
  File "test.py", lineno 33
    b = Branch(self)  # this isn't a useful line to warn about

So another question is: how does partial.__call__ get skipped in the stack frames and can I replicate that for a regular method?

EDIT: I do know about PYTHONTRACEMALLOC=2 to get the additional frame for all resource warnings, but it would be ideal if I could skip it by default just like the stacklevel argument.

:wave: Hi Min! Just ran into this randomly while searching for an example to add ResourceWarning to some of my own code :slight_smile:

I’m pretty sure there is no way to skip stack frames in the tracemalloc output, currently. I guess the only way would be to use the tracemalloc module to fetch the stack trace and then format it manually, which seems like more trouble than it’s worth.

The behavior you saw with partial is because partial is implemented in C, and it doesn’t create a Python-level stack frame at all. In principle you could do this too by implementing your allocation method in C, but again, seems like more trouble than it’s worth.