📒 Allow executable code/output to be inside a figure

.changeset/nasty-carrots-attack.md

@@ -0,0 +1,6 @@
+---
+'@myst-theme/jupyter': patch
+'@myst-theme/site': patch
+---
+
+Allow executable code/output to be inside a figure

packages/jupyter/src/execute/index.ts

@@ -4,3 +4,4 @@ export * from './provider.js';
 export * from './selectors.js';
 export * from './types.js';
 export * from './busy.js';
+export * from './utils.js';

packages/jupyter/src/execute/utils.ts

@@ -2,6 +2,20 @@ import type { GenericParent } from 'myst-common';
 import type { Config, IRenderMimeRegistry, ThebeCore } from 'thebe-core';
 import type { IdKeyMap, IdKeyMapTarget } from './types.js';
 
+/**
+ * Return executable code/output from block or single figure inside block
+ */
+export function executableNodesFromBlock(block: GenericParent) {
+  if (!block || block.type !== 'block') return;
+  let target = block;
+  if (block.children && block.children.length === 1 && block.children[0].type === 'container') {
+    target = block.children[0] as GenericParent;
+  }
+  if (target.children && target.children.length >= 2 && target.children[0].type === 'code') {
+    return { codeCell: target.children[0], output: target.children[1] };
+  }
+}
+
 /**
  * Use the mdast to create a ThebeNotebook from the mdast tree of a notebook.
  * This is intended to be used to create an independent ThebeNotebook instance
@@ -40,8 +54,9 @@ export function notebookFromMdast(
 
   notebook.cells = (mdast.children as GenericParent[]).map((block: GenericParent) => {
     if (block.type !== 'block') console.warn(`Unexpected block type ${block.type}`);
-    if (block.children.length == 2 && block.children[0].type === 'code') {
-      const [codeCell, output] = block.children;
+    const executableNodes = executableNodesFromBlock(block);
+    if (executableNodes) {
+      const { codeCell, output } = executableNodes;
 
       // use the block.key to identify the cell but maintain a mapping
       // to allow code or output keys to look up cells and refs and idenifity

packages/site/src/components/ContentBlocks.tsx

@@ -3,21 +3,15 @@ import { SourceFileKind } from 'myst-spec-ext';
 import type { GenericParent } from 'myst-common';
 import classNames from 'classnames';
 import {
+  executableNodesFromBlock,
   NotebookClearCell,
   NotebookRunCell,
   NotebookRunCellSpinnerOnly,
 } from '@myst-theme/jupyter';
 import { useGridSystemProvider } from '@myst-theme/providers';
 
 function isACodeCell(node: GenericParent) {
-  return (
-    node &&
-    node.type === 'block' &&
-    node.children &&
-    node.children?.length === 2 &&
-    node.children[0].type === 'code' &&
-    node.children[1].type === 'output'
-  );
+  return !!executableNodesFromBlock(node);
 }
 
 function Block({