inline doctest properly

docs/_build/scripts/macro.py

@@ -131,6 +131,23 @@ def code_tab(
 
 
 def define_env(env):
+    @env.macro
+    def code_header(language: str, section: str = [], api_functions: List[str] = []) -> str:
+        language_info = LANGUAGES[language]
+
+        language = language_info["code_name"]
+
+        # Create feature flags
+        feature_flags_links = create_feature_flag_links(language, api_functions)
+
+        # Create API Links if they are defined in the YAML
+        api_functions = [
+            link for f in api_functions if (link := create_api_function_link(language, f))
+        ]
+        language_headers = " ·".join(api_functions + feature_flags_links)
+        return f"""=== \":fontawesome-brands-{language_info['icon_name']}: {language_info['display_name']}\"
+    {language_headers}"""
+
     @env.macro
     def code_block(
         path: str, section: str = None, api_functions: List[str] = None

docs/src/python/user-guide/lazy/gpu.py

@@ -1,38 +1,28 @@
 # --8<-- [start:setup]
-# --8<-- [start:simple]
 import polars as pl
 
 df = pl.LazyFrame({"a": [1.242, 1.535]})
 
 q = df.select(pl.col("a").round(1))
 # --8<-- [end:setup]
 
-result = q.collect(engine="gpu")
-print(result)
-# --8<-- [end:simple]
-
 # Avoiding requiring the GPU engine for doc build output
 # --8<-- [start:simple-result]
 result = q.collect()
 print(result)
 # --8<-- [end:simple-result]
 
 
-# --8<-- [start:engine]
 # --8<-- [start:engine-setup]
 q = df.select((pl.col("a") ** 4))
 
 # --8<-- [end:engine-setup]
-result = q.collect(engine=pl.GPUEngine(device=1))
-print(result)
-# --8<-- [end:engine]
 
 # --8<-- [start:engine-result]
 result = q.collect()
 print(result)
 # --8<-- [end:engine-result]
 
-# --8<-- [start:fallback-warning]
 # --8<-- [start:fallback-setup]
 df = pl.LazyFrame(
     {
@@ -45,12 +35,6 @@
 
 # --8<-- [end:fallback-setup]
 
-with pl.Config() as cfg:
-    cfg.set_verbose(True)
-    result = q.collect(engine="gpu")
-
-print(result)
-# --8<-- [end:fallback-warning]
 # --8<-- [start:fallback-result]
 print(
     "PerformanceWarning: Query execution with GPU not supported, reason: \n"

docs/user-guide/gpu-support.md

@@ -30,7 +30,13 @@ You can install the GPU backend for Polars with a feature flag as part of a norm
 
 Having built a query using the lazy API [as normal](lazy/index.md), GPU-enabled execution is requested by running `.collect(engine="gpu")` instead of `.collect()`.
 
-{{code_block('user-guide/lazy/gpu', 'simple', ['collect'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:setup"
+
+result = q.collect(engine="gpu")
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:setup"
@@ -39,7 +45,12 @@ Having built a query using the lazy API [as normal](lazy/index.md), GPU-enabled
 
 For more detailed control over the execution, for example to specify which GPU to use on a multi-GPU node, we can provide a `GPUEngine` object. By default, the GPU engine will use a configuration applicable to most use cases.
 
-{{code_block('user-guide/lazy/gpu', 'engine', ['GPUEngine'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:engine-setup"
+result = q.collect(engine=pl.GPUEngine(device=1))
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:engine-setup"
@@ -87,16 +98,32 @@ The release of the GPU engine in Open Beta implies that we expect things to work
 
 When running in verbose mode, any queries that cannot execute on the GPU will issue a `PerformanceWarning`:
 
-{{code_block('user-guide/lazy/gpu', 'fallback-warning', ['Config'])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:fallback-setup"
+
+with pl.Config() as cfg:
+    cfg.set_verbose(True)
+    result = q.collect(engine="gpu")
+
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:fallback-setup"
---8<-- "python/user-guide/lazy/gpu.py:fallback-result"
+print(
+    "PerformanceWarning: Query execution with GPU not supported, reason: \n"
+    "<class 'NotImplementedError'>: Grouped rolling window not implemented"
+)
+print("# some details elided")
+print()
+print(q.collect())
 ```
 
 To disable fallback, and have the GPU engine raise an exception if a query is unsupported, we can pass an appropriately configured `GPUEngine` object:
 
-```python exec="on" result="text" session="user-guide/lazy"
+{{ code_header("python", [], []) }}
+```python
 q.collect(engine=pl.GPUEngine(raise_on_fail=True))
 ```
 

docs/user-guide/lazy/gpu.md

@@ -4,7 +4,13 @@ Polars provides an in-memory, GPU-accelerated execution engine for the Lazy API
 
 If you install Polars with the [GPU feature flag](../installation.md), you can trigger GPU-based execution by running `.collect(engine="gpu")` instead of `.collect()`.
 
-{{code_block('user-guide/lazy/gpu', 'simple', [])}}
+{{ code_header("python", [], []) }}
+```python
+--8<-- "python/user-guide/lazy/gpu.py:setup"
+
+result = q.collect(engine="gpu")
+print(result)
+```
 
 ```python exec="on" result="text" session="user-guide/lazy"
 --8<-- "python/user-guide/lazy/gpu.py:setup"