diff --git a/custom_components/pyscript/eval.py b/custom_components/pyscript/eval.py index 8a9df38..6e0990e 100644 --- a/custom_components/pyscript/eval.py +++ b/custom_components/pyscript/eval.py @@ -1280,43 +1280,50 @@ async def ast_raise(self, arg): raise # pylint: disable=misplaced-bare-raise async def ast_with(self, arg, async_attr=""): - """Execute with statement.""" - hit_except = False - ctx_list = [] + """ + Execute with statement. + + Several items behave like nested with statements: each manager is + evaluated and entered in turn, only managers that were entered are + exited, and they are exited innermost first. An exit that returns true + suppresses the error, so the outer exits see no error; an exit that + raises replaces the error for the outer exits. + """ + entered = [] val = None + error = None enter_attr = f"__{async_attr}enter__" exit_attr = f"__{async_attr}exit__" try: for item in arg.items: manager = await self.aeval(item.context_expr) - ctx_list.append( - { - "manager": manager, - "enter": getattr(type(manager), enter_attr), - "exit": getattr(type(manager), exit_attr), - "target": item.optional_vars, - } - ) - for ctx in ctx_list: - value = await self.call_func(ctx["enter"], enter_attr, ctx["manager"]) - if ctx["target"]: - await self.recurse_assign(ctx["target"], value) + enter_func = getattr(type(manager), enter_attr) + exit_func = getattr(type(manager), exit_attr) + value = await self.call_func(enter_func, enter_attr, manager) + entered.append((manager, exit_func)) + if item.optional_vars: + await self.recurse_assign(item.optional_vars, value) for arg1 in arg.body: val = await self.aeval(arg1) if isinstance(val, EvalStopFlow): break - except Exception: - hit_except = True - exit_ok = True - for ctx in reversed(ctx_list): - ret = await self.call_func(ctx["exit"], exit_attr, ctx["manager"], *sys.exc_info()) - exit_ok = exit_ok and ret - if not exit_ok: - raise - finally: - if not hit_except: - for ctx in reversed(ctx_list): - await self.call_func(ctx["exit"], exit_attr, ctx["manager"], None, None, None) + except BaseException as exc: # pylint: disable=broad-exception-caught + # also CancelledError, which the exits must see like in Python + error = exc + for manager, exit_func in reversed(entered): + try: + if error is None: + await self.call_func(exit_func, exit_attr, manager, None, None, None) + elif await self.call_func( + exit_func, exit_attr, manager, type(error), error, error.__traceback__ + ): + error = None + except BaseException as exc: # pylint: disable=broad-exception-caught + if error is not None and exc is not error and exc.__context__ is None: + exc.__context__ = error + error = exc + if error is not None: + raise error return val async def ast_asyncwith(self, arg): diff --git a/tests/test_unit_eval.py b/tests/test_unit_eval.py index 27ee386..31e3c26 100644 --- a/tests/test_unit_eval.py +++ b/tests/test_unit_eval.py @@ -1,6 +1,9 @@ """Unit tests for Python interpreter.""" +import ast as ast_module +import contextlib from dataclasses import dataclass +import types from types import ModuleType import pytest @@ -2185,3 +2188,123 @@ def test_reraise(): await Function.waiter_sync() await Function.waiter_stop() await Function.reaper_stop() + + +WITH_SEMANTICS_SOURCE = """ +class Ctx: + def __init__(self, msgs, name, fail_enter=False, suppress=False, fail_exit=False): + self.msgs = msgs + self.name = name + self.fail_enter = fail_enter + self.suppress = suppress + self.fail_exit = fail_exit + + def __enter__(self): + self.msgs.append(f"enter {self.name}") + if self.fail_enter: + raise ValueError(f"enter {self.name}") + return self.name + + def __exit__(self, exc_type, exc, tb): + self.msgs.append(f"exit {self.name} {exc_type.__name__ if exc_type else None}") + if self.fail_exit: + raise KeyError(f"exit {self.name}") + return self.suppress + +class ACtx: + def __init__(self, msgs, name, **kwargs): + self.ctx = Ctx(msgs, name, **kwargs) + + async def __aenter__(self): + return self.ctx.__enter__() + + async def __aexit__(self, exc_type, exc, tb): + return self.ctx.__exit__(exc_type, exc, tb) + +def make(msgs, name, **kwargs): + msgs.append(f"make {name}") + return Ctx(msgs, name, **kwargs) + +def amake(msgs, name, **kwargs): + msgs.append(f"make {name}") + return ACtx(msgs, name, **kwargs) + +def describe(exc): + context = type(exc.__context__).__name__ if exc.__context__ else None + return f"raised {type(exc).__name__} {exc} context {context}" + +results = [] +cases = [ + {"b": {"fail_enter": True}}, + {"body": True, "b": {"suppress": True}}, + {"body": True, "a": {"suppress": True}}, + {"body": True, "b": {"fail_exit": True}}, + {"b": {"fail_exit": True}}, + {"body": True, "a": {"fail_exit": True}, "b": {"suppress": True}}, + {}, +] +for case in cases: + msgs = [] + try: + with make(msgs, "a", **case.get("a", {})) as a, make(msgs, "b", **case.get("b", {})) as b: + msgs.append(f"body {a} {b}") + if case.get("body"): + raise ZeroDivisionError("body") + msgs.append("no error") + except Exception as exc: + msgs.append(describe(exc)) + results.append(msgs) + + msgs = [] + try: + async with amake(msgs, "a", **case.get("a", {})) as a, amake(msgs, "b", **case.get("b", {})) as b: + msgs.append(f"body {a} {b}") + if case.get("body"): + raise ZeroDivisionError("body") + msgs.append("no error") + except Exception as exc: + msgs.append(describe(exc)) + results.append(msgs) + +msgs = [] +try: + async with failing_acm(): + msgs.append("body") +except Exception as exc: + msgs.append(describe(exc)) +results.append(msgs) +""" + + +@contextlib.asynccontextmanager +async def failing_acm(): + """Fail on enter, like a generator based manager whose connection is refused.""" + raise ConnectionRefusedError("refused") + yield # pylint: disable=unreachable + + +async def test_with_matches_python(hass): + """Test that with and async with behave like Python, compared with CPython itself.""" + hass.data[DOMAIN] = {CONFIG_ENTRY: MockConfigEntry(domain=DOMAIN, data={CONF_ALLOW_ALL_IMPORTS: True})} + Function.init(hass) + State.init(hass) + State.register_functions() + TrigTime.init(hass) + + expected_ns = {"failing_acm": failing_acm} + code = compile(WITH_SEMANTICS_SOURCE, "test", "exec", flags=ast_module.PyCF_ALLOW_TOP_LEVEL_AWAIT) + # code compiled with top level await runs as a coroutine function, like the asyncio REPL + await types.FunctionType(code, expected_ns)() + expected = expected_ns["results"] + + global_ctx = GlobalContext( + "test", global_sym_table={"failing_acm": failing_acm}, manager=GlobalContextMgr + ) + ast = AstEval("test", global_ctx=global_ctx) + Function.install_ast_funcs(ast) + ast.parse(WITH_SEMANTICS_SOURCE + "results\n") + assert await ast.eval() == expected + + await Function.waiter_sync() + await Function.waiter_stop() + await Function.reaper_stop()