Exit only the context managers that were entered in with and async with

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):

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()