From 2e0f878bdd73494012fa58ae612e3ee684210a41 Mon Sep 17 00:00:00 2001
From: Franklin Koch <franklinwkoch@gmail.com>
Date: Fri, 17 Nov 2023 02:53:20 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=93=92=20Allow=20executable=20code/output?=
 =?UTF-8?q?=20to=20be=20inside=20a=20figure?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .changeset/nasty-carrots-attack.md            |  6 ++++++
 packages/jupyter/src/execute/index.ts         |  1 +
 packages/jupyter/src/execute/utils.ts         | 19 +++++++++++++++++--
 .../site/src/components/ContentBlocks.tsx     | 10 ++--------
 4 files changed, 26 insertions(+), 10 deletions(-)
 create mode 100644 .changeset/nasty-carrots-attack.md

diff --git a/.changeset/nasty-carrots-attack.md b/.changeset/nasty-carrots-attack.md
new file mode 100644
index 000000000..056d3c53b
--- /dev/null
+++ b/.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
diff --git a/packages/jupyter/src/execute/index.ts b/packages/jupyter/src/execute/index.ts
index 7677233f2..b35f807f1 100644
--- a/packages/jupyter/src/execute/index.ts
+++ b/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';
diff --git a/packages/jupyter/src/execute/utils.ts b/packages/jupyter/src/execute/utils.ts
index 4bbea50ed..65432d9e8 100644
--- a/packages/jupyter/src/execute/utils.ts
+++ b/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
diff --git a/packages/site/src/components/ContentBlocks.tsx b/packages/site/src/components/ContentBlocks.tsx
index 22795a178..0702e2e32 100644
--- a/packages/site/src/components/ContentBlocks.tsx
+++ b/packages/site/src/components/ContentBlocks.tsx
@@ -3,6 +3,7 @@ import { SourceFileKind } from 'myst-spec-ext';
 import type { GenericParent } from 'myst-common';
 import classNames from 'classnames';
 import {
+  executableNodesFromBlock,
   NotebookClearCell,
   NotebookRunCell,
   NotebookRunCellSpinnerOnly,
@@ -10,14 +11,7 @@ import {
 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({