add support for async galleries

docs/examples/plot_10_async.py

@@ -0,0 +1,48 @@
+"""
+# Async support
+"""
+
+import asyncio
+import time
+
+start = time.time()
+await asyncio.sleep(1)
+stop = time.time()
+f"I waited for {stop - start} seconds!"
+
+
+# %%
+# ## Async Iterator
+
+class AsyncIterator:
+    async def __aiter__(self):
+        for chunk in "I'm an async iterator!".split():
+            yield chunk
+
+
+async for chunk in AsyncIterator():
+    print(chunk, end=" ")
+
+# %%
+# ## Async comprehensions
+
+" ".join([chunk async for chunk in AsyncIterator()])
+
+# %%
+# ## Async content manager
+
+
+class AsyncContextManager:
+    async def __aenter__(self):
+        print("Entering ...")
+        return self
+
+    async def __aexit__(self, *exc_info):
+        print("Exiting ...")
+
+    def __str__(self):
+        return "I'm an async context manager!"
+
+
+async with AsyncContextManager() as acm:
+    print(acm)

src/mkdocs_gallery/gen_single.py

@@ -27,7 +27,7 @@
 from io import StringIO
 from pathlib import Path
 from shutil import copyfile
-from textwrap import indent
+from textwrap import indent, dedent
 from time import time
 from typing import List, Set, Tuple
 
@@ -739,6 +739,66 @@ def _reset_cwd_syspath(cwd, path_to_remove):
     os.chdir(cwd)
 
 
+def _parse_code(bcontent, src_file, *, compiler_flags):
+    code_ast = compile(bcontent, src_file, "exec", compiler_flags | ast.PyCF_ONLY_AST, dont_inherit=1)
+    if _needs_async_handling(bcontent, compiler_flags=compiler_flags):
+        code_ast = _apply_async_handling(code_ast, compiler_flags=compiler_flags)
+    return code_ast
+
+
+def _needs_async_handling(bcontent, *, compiler_flags) -> bool:
+    try:
+        compile(bcontent, "<_needs_async_handling()>", "exec", compiler_flags, dont_inherit=1)
+    except SyntaxError as error:
+        # FIXME
+        #  bool(re.match(r"'(await|async for|async with)' outside( async)? function", str(error)))
+        #  asynchronous comprehension outside of an asynchronous function
+        return "async" in str(error)
+    except Exception:
+        return False
+    else:
+        return False
+
+
+def _apply_async_handling(code_ast, *, compiler_flags):
+    async_handling = compile(
+        dedent(
+            """
+                async def __async_wrapper__():
+                    # original AST goes here
+                    return locals()
+                import asyncio as __asyncio__
+                __async_wrapper_locals__ = __asyncio__.run(__async_wrapper__())
+                __async_wrapper_result__ = __async_wrapper_locals__.pop("__async_wrapper_result__", None)
+                globals().update(__async_wrapper_locals__)
+                __async_wrapper_result__
+                """
+        ),
+        "<_apply_async_handling()>",
+        "exec",
+        compiler_flags | ast.PyCF_ONLY_AST,
+        dont_inherit=1,
+    )
+
+    *original_body, last_node = code_ast.body
+    if isinstance(last_node, ast.Expr):
+        # FIXME
+        #  ast.Assign(targets=[ast.Name(id="__async_wrapper_result__", ctx=ast.Store())], value=last_node)
+        last_node = compile(
+            f"__async_wrapper_result__ = ({ast.unparse(last_node)})",
+            "<_apply_async_handling()>",
+            "exec",
+            compiler_flags | ast.PyCF_ONLY_AST,
+            dont_inherit=1,
+        ).body[0]
+    original_body.append(last_node)
+
+    async_wrapper = async_handling.body[0]
+    async_wrapper.body = [*original_body, *async_wrapper.body]
+
+    return ast.fix_missing_locations(async_handling)
+
+
 def execute_code_block(compiler, block, script: GalleryScript):
     """Execute the code block of the example file.
 
@@ -788,9 +848,7 @@ def execute_code_block(compiler, block, script: GalleryScript):
 
     try:
         ast_Module = _ast_module()
-        code_ast = ast_Module([bcontent])
-        flags = ast.PyCF_ONLY_AST | compiler.flags
-        code_ast = compile(bcontent, src_file, "exec", flags, dont_inherit=1)
+        code_ast = _parse_code(bcontent, src_file, compiler_flags=compiler.flags)
         ast.increment_lineno(code_ast, lineno - 1)
 
         is_last_expr, mem_max = _exec_and_get_memory(compiler, ast_Module, code_ast, script=script)