A while back we added tests to make sure we weren't leaking hiwire keys. As a special case of this, we also were testing that we weren't leaking JsProxies. In the future, I would like to add tests to show that we aren't leaking PyProxies either. Currently we leak quite a lot of them.

Before we actually add any tests to show that we aren't leaking proxies, I'd like to rework the model of the PyProxy lifetime. The planned changes should make it far easier to actually avoid these leaks while still using a reasonably natural coding style.

This change has no user-facing effect: the tracing is disabled by default and the API to enable it is private and always will be.
Add support for pyproxy allocation tracing. A while back we added tests to make sure we weren't leaking hiwire keys. As a special case of this, we also were testing that we weren't leaking JsProxies. In the future, I would like to add tests to show that we aren't leaking PyProxies either. Currently we leak quite a lot of them.

Before we actually add any tests to show that we aren't leaking proxies, I'd like to rework the model of the PyProxy lifetime. The planned changes should make it far easier to actually avoid these leaks while still using a reasonably natural coding style.

This change has no user-facing effect: the tracing is disabled by default and the API to enable it is private and always will be.
import __main__
import sys
import gc

from ._core import JsProxy
from ._importhooks import jsfinder


def save_state() -> dict:
    """Record the current global state.

    This includes which Javascript packages are loaded and the global scope in
    ``__main__.__dict__``. Many loaded modules might have global state, but
    there is no general way to track it and we don't try to.
    """
    loaded_js_modules = {}
    for [key, value] in sys.modules.items():
        if isinstance(value, JsProxy):
            loaded_js_modules[key] = value

    return dict(
        globals=dict(__main__.__dict__),
        js_modules=dict(jsfinder.jsproxies),
        loaded_js_modules=loaded_js_modules,
    )


def restore_state(state: dict):
    """Restore the global state to a snapshot. The argument ``state`` should
    come from ``save_state``"""
    __main__.__dict__.clear()
    __main__.__dict__.update(state["globals"])

    jsfinder.jsproxies = state["js_modules"]
    loaded_js_modules = state["loaded_js_modules"]
    for [key, value] in list(sys.modules.items()):
        if isinstance(value, JsProxy) and key not in loaded_js_modules:
            del sys.modules[key]
    sys.modules.update(loaded_js_modules)

    sys.last_type = None
    sys.last_value = None
    sys.last_traceback = None

    return gc.collect(2)


__all__ = ["save_state", "restore_state"]

import __main__
import sys
import gc

from ._core import JsProxy
from ._importhooks import jsfinder


def save_state() -> dict:
    """Record the current global state.

    This includes which Javascript packages are loaded and the global scope in
    ``__main__.__dict__``. Many loaded modules might have global state, but
    there is no general way to track it and we don't try to.
    """
    loaded_js_modules = {}
    for [key, value] in sys.modules.items():
        if isinstance(value, JsProxy):
            loaded_js_modules[key] = value

    return dict(
        globals=dict(__main__.__dict__),
        js_modules=dict(jsfinder.jsproxies),
        loaded_js_modules=loaded_js_modules,
    )


def restore_state(state: dict):
    """Restore the global state to a snapshot. The argument ``state`` should
    come from ``save_state``"""
    __main__.__dict__.clear()
    __main__.__dict__.update(state["globals"])

    jsfinder.jsproxies = state["js_modules"]
    loaded_js_modules = state["loaded_js_modules"]
    for [key, value] in list(sys.modules.items()):
        if isinstance(value, JsProxy) and key not in loaded_js_modules:
            del sys.modules[key]
    sys.modules.update(loaded_js_modules)

    sys.last_type = None
    sys.last_value = None
    sys.last_traceback = None

    return gc.collect(2)


__all__ = ["save_state", "restore_state"]
