feat: add top-level await support for async code blocks (#62)

* feat: add top-level await support for async code blocks

Collapse ast.parse() + compile(tree) into a single compile(source) call
with PyCF_ALLOW_TOP_LEVEL_AWAIT flag. Detect coroutine code objects via
CO_COROUTINE and run them with asyncio.run(), keeping sync path unchanged.

This allows doc authors to write bare `await`, `async for`, and
`async with` in code blocks without wrapping in asyncio.run() boilerplate.

* fix: use pytest-asyncio's event loop for top-level await code blocks

asyncio.run() creates a new event loop per call, breaking async fixtures
that bind resources to pytest-asyncio's managed loop. Now fetches the
_function_scoped_runner from pytest-asyncio (EAFP) and runs coroutines
on the shared loop. Falls back to asyncio.run() when pytest-asyncio is
absent or the fixture lookup fails.

* fix: require pytest-asyncio for async code blocks instead of falling back to asyncio.run()

asyncio.run() creates a new event loop, risking silent buggy behavior
when the pytest-asyncio runner can't be found. Now raises a clear
RuntimeError directing the user to install pytest-asyncio, mirroring
how pytest itself handles async test functions without a plugin.

Author: jjovalle99

src/pytest_markdown_docs/_runners.py

@@ -1,5 +1,6 @@
 import abc
 import ast
+import inspect
 import traceback
 import typing
 from abc import abstractmethod
@@ -50,22 +51,29 @@ def decorator(r: RUNNER_TYPE) -> RUNNER_TYPE:
 
 @register_runner(default=True)
 class DefaultRunner(_Runner):
-    def runtest(self, test: FenceTestDefinition, args):
+    def runtest(self, test: FenceTestDefinition, args, *, asyncio_runner=None):
         try:
-            tree = ast.parse(test.source, filename=test.source_path)
-        except SyntaxError:
-            raise
-
-        try:
-            # if we don't compile the code, it seems we get name lookup errors
-            # for functions etc. when doing cross-calls across inline functions
             compiled = compile(
-                tree, filename=test.source_path, mode="exec", dont_inherit=True
+                test.source,
+                filename=test.source_path,
+                mode="exec",
+                flags=ast.PyCF_ALLOW_TOP_LEVEL_AWAIT,
+                dont_inherit=True,
             )
         except SyntaxError:
             raise
 
-        exec(compiled, args)
+        if compiled.co_flags & inspect.CO_COROUTINE:
+            if asyncio_runner is None:
+                raise RuntimeError(
+                    "Top-level async code in markdown code blocks is not natively supported.\n"
+                    "You need to install pytest-asyncio to run async code blocks:\n"
+                    "  pip install pytest-asyncio"
+                )
+            coro = eval(compiled, args)
+            asyncio_runner.run(coro)
+        else:
+            exec(compiled, args)
 
     def repr_failure(
         self,

src/pytest_markdown_docs/plugin.py

@@ -52,6 +52,14 @@ def get_docstring_start_line(obj) -> typing.Optional[int]:
     return None  # Docstring not found in source
 
 
+def _get_asyncio_runner(fixture_request):
+    """Try to fetch pytest-asyncio's event loop runner for shared-loop execution."""
+    try:
+        return fixture_request.getfixturevalue("_function_scoped_runner")
+    except Exception:
+        return None
+
+
 class MarkdownInlinePythonItem(pytest.Item):
     def __init__(
         self,
@@ -109,8 +117,17 @@ def runtest(self):
             try:
                 # this ensures that pytest's stdout/stderr capture works during the test:
                 capman = self.config.pluginmanager.getplugin("capturemanager")
+                asyncio_runner = _get_asyncio_runner(self.fixture_request)
                 with capman.global_and_fixture_disabled():
-                    self.runner.runtest(self.test_definition, all_globals)
+                    try:
+                        self.runner.runtest(
+                            self.test_definition,
+                            all_globals,
+                            asyncio_runner=asyncio_runner,
+                        )
+                    except TypeError:
+                        # Custom runner doesn't accept asyncio_runner kwarg
+                        self.runner.runtest(self.test_definition, all_globals)
 
                 # Success - test passed
                 if attempt > 0:

tests/plugin_test.py

@@ -976,3 +976,136 @@ def sync_value():
     )
     result = testdir.runpytest("--markdown-docs", "-v")
     result.assert_outcomes(passed=1)
+
+
+# ============================================================================
+# Top-level await tests
+# ============================================================================
+
+
+def test_top_level_await(testdir):
+    """Test that bare await at top level works in code blocks."""
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python
+import asyncio
+await asyncio.sleep(0)
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs")
+    result.assert_outcomes(passed=1)
+
+
+def test_top_level_async_for(testdir):
+    """Test that async for at top level works in code blocks."""
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python
+import asyncio
+
+async def arange(n):
+    for i in range(n):
+        yield i
+
+results = []
+async for i in arange(3):
+    results.append(i)
+assert results == [0, 1, 2]
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs")
+    result.assert_outcomes(passed=1)
+
+
+def test_top_level_async_with(testdir):
+    """Test that async with at top level works in code blocks."""
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python
+import contextlib
+
+@contextlib.asynccontextmanager
+async def async_ctx():
+    yield "hello"
+
+async with async_ctx() as val:
+    assert val == "hello"
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs")
+    result.assert_outcomes(passed=1)
+
+
+def test_continuation_with_await(testdir):
+    """Test that continuation blocks can use await with state from sync blocks."""
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python
+value = 42
+```
+
+```python continuation
+import asyncio
+await asyncio.sleep(0)
+assert value == 42
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs")
+    result.assert_outcomes(passed=2)
+
+
+def test_async_fixture_shares_event_loop(testdir):
+    """Async fixture and top-level await code block must share the same event loop."""
+    testdir.makeini("[pytest]\nasyncio_mode = auto\n")
+    testdir.makeconftest(
+        """
+import asyncio
+import pytest
+
+@pytest.fixture
+async def loop_id():
+    return id(asyncio.get_running_loop())
+"""
+    )
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python fixture:loop_id
+import asyncio
+await asyncio.sleep(0)
+assert loop_id == id(asyncio.get_running_loop())
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs", "-v")
+    result.assert_outcomes(passed=1)
+
+
+def test_top_level_await_requires_pytest_asyncio(testdir):
+    """Async code blocks should error clearly when pytest-asyncio is not installed."""
+    testdir.makeconftest(
+        """
+import pytest_markdown_docs.plugin as _plugin
+_plugin._get_asyncio_runner = lambda *a, **kw: None
+"""
+    )
+    testdir.makefile(
+        ".md",
+        test_file="""
+```python
+import asyncio
+await asyncio.sleep(0)
+```
+""",
+    )
+    result = testdir.runpytest("--markdown-docs")
+    result.assert_outcomes(failed=1)
+    result.stdout.fnmatch_lines(["*pytest-asyncio*"])