Fix bug publish screenshot (#762)

[Preview this branch](https://cosmos-explorer-preview.azurewebsites.net/pull/762?feature.someFeatureFlagYouMightNeed=true)

The main change in this PR fixes the snapshot functionality in the Publish pane-related components. Because the code cell outputs are now rendered in their own iframes for security reasons, a single snapshot of the notebook is no longer possible: each cell output takes its own snapshot and the snapshots are collated on the main notebook snapshot.
- Move the snapshot functionality to notebook components: this removes the reference of the notebook DOM node that we must pass to the Publish pane via explorer.
- Add slice in the state and actions in notebook redux for notebook snapshot requests and result
- Add post robot message to take snapshots and receive results
- Add logic in `NotebookRenderer` to wait for all output snapshots done before taking the main one collating.
- Use `zustand` to share snapshot between Redux world and React world. This solves the issue of keeping the `PanelContainer` component generic, while being able to update its children (`PublishPanel` component) with the new snapshot.

Additional changes:
- Add `local()` in `@font-face` to check if font is already installed before downloading the font (must be done for Safari, but not Edge/Chrome)
- Add "Export output to image" menu item in notebook cell, since each cell output can take its own snapshot (which can be downloaded)
![image](https://user-images.githubusercontent.com/21954022/117454706-b5f16600-af46-11eb-8535-6bf99f3d9170.png)

src/Explorer/Notebook/NotebookComponent/NotebookComponentAdapter.tsx

@@ -1,14 +1,11 @@
-import * as React from "react";
-
-import { ReactAdapter } from "../../../Bindings/ReactBindingHandler";
-import { NotebookClientV2 } from "../NotebookClientV2";
-
 // Vendor modules
 import { actions, createContentRef, createKernelRef, selectors } from "@nteract/core";
-import VirtualCommandBarComponent from "./VirtualCommandBarComponent";
+import * as React from "react";
+import { ReactAdapter } from "../../../Bindings/ReactBindingHandler";
+import { NotebookClientV2 } from "../NotebookClientV2";
 import { NotebookContentItem } from "../NotebookContentItem";
 import { NotebookComponentBootstrapper } from "./NotebookComponentBootstrapper";
-import { CdbAppState } from "./types";
+import VirtualCommandBarComponent from "./VirtualCommandBarComponent";
 
 export interface NotebookComponentAdapterOptions {
   contentItem: NotebookContentItem;
@@ -19,7 +16,6 @@ export interface NotebookComponentAdapterOptions {
 
 export class NotebookComponentAdapter extends NotebookComponentBootstrapper implements ReactAdapter {
   private onUpdateKernelInfo: () => void;
-  public getNotebookParentElement: () => HTMLElement;
   public parameters: any;
 
   constructor(options: NotebookComponentAdapterOptions) {
@@ -46,11 +42,6 @@ export class NotebookComponentAdapter extends NotebookComponentBootstrapper impl
         })
       );
     }
-
-    this.getNotebookParentElement = () => {
-      const cdbAppState = this.getStore().getState() as CdbAppState;
-      return cdbAppState.cdb.currentNotebookParentElements.get(this.contentRef);
-    };
   }
 
   protected renderExtraComponent = (): JSX.Element => {

src/Explorer/Notebook/NotebookComponent/NotebookComponentBootstrapper.tsx

@@ -1,35 +1,29 @@
-import * as React from "react";
-
-import { NotebookComponent } from "./NotebookComponent";
-import { NotebookClientV2 } from "../NotebookClientV2";
-import { NotebookUtil } from "../NotebookUtil";
-
+import { CellId, CellType, ImmutableNotebook } from "@nteract/commutable";
 // Vendor modules
 import {
   actions,
   AppState,
-  createKernelRef,
-  DocumentRecordProps,
   ContentRef,
+  DocumentRecordProps,
   KernelRef,
   NotebookContentRecord,
   selectors,
 } from "@nteract/core";
-import * as Immutable from "immutable";
-import { Provider } from "react-redux";
-import { CellType, CellId, ImmutableNotebook } from "@nteract/commutable";
-import { Store, AnyAction } from "redux";
-
-import "./NotebookComponent.less";
-
-import "codemirror/addon/hint/show-hint.css";
-import "codemirror/lib/codemirror.css";
 import "@nteract/styles/editor-overrides.css";
 import "@nteract/styles/global-variables.css";
+import "codemirror/addon/hint/show-hint.css";
+import "codemirror/lib/codemirror.css";
+import * as Immutable from "immutable";
+import * as React from "react";
+import { Provider } from "react-redux";
 import "react-table/react-table.css";
-
-import * as CdbActions from "./actions";
+import { AnyAction, Store } from "redux";
+import { NotebookClientV2 } from "../NotebookClientV2";
+import { NotebookUtil } from "../NotebookUtil";
 import * as NteractUtil from "../NTeractUtil";
+import * as CdbActions from "./actions";
+import { NotebookComponent } from "./NotebookComponent";
+import "./NotebookComponent.less";
 
 export interface NotebookComponentBootstrapperOptions {
   notebookClient: NotebookClientV2;
@@ -37,7 +31,7 @@ export interface NotebookComponentBootstrapperOptions {
 }
 
 export class NotebookComponentBootstrapper {
-  protected contentRef: ContentRef;
+  public contentRef: ContentRef;
   protected renderExtraComponent: () => JSX.Element;
 
   private notebookClient: NotebookClientV2;

src/Explorer/Notebook/NotebookComponent/actions.ts

@@ -1,6 +1,7 @@
 import { CellId } from "@nteract/commutable";
 import { ContentRef } from "@nteract/core";
 import { Action } from "../../../Shared/Telemetry/TelemetryConstants";
+import { SnapshotFragment, SnapshotRequest } from "./types";
 
 export const CLOSE_NOTEBOOK = "CLOSE_NOTEBOOK";
 export interface CloseNotebookAction {
@@ -85,21 +86,68 @@ export const traceNotebookTelemetry = (payload: {
   };
 };
 
-export const UPDATE_NOTEBOOK_PARENT_DOM_ELTS = "UPDATE_NOTEBOOK_PARENT_DOM_ELTS";
-export interface UpdateNotebookParentDomEltAction {
-  type: "UPDATE_NOTEBOOK_PARENT_DOM_ELTS";
+export const STORE_CELL_OUTPUT_SNAPSHOT = "STORE_CELL_OUTPUT_SNAPSHOT";
+export interface StoreCellOutputSnapshotAction {
+  type: "STORE_CELL_OUTPUT_SNAPSHOT";
   payload: {
-    contentRef: ContentRef;
-    parentElt: HTMLElement;
+    cellId: string;
+    snapshot: SnapshotFragment;
   };
 }
 
-export const UpdateNotebookParentDomElt = (payload: {
-  contentRef: ContentRef;
-  parentElt: HTMLElement;
-}): UpdateNotebookParentDomEltAction => {
+export const storeCellOutputSnapshot = (payload: {
+  cellId: string;
+  snapshot: SnapshotFragment;
+}): StoreCellOutputSnapshotAction => {
   return {
-    type: UPDATE_NOTEBOOK_PARENT_DOM_ELTS,
+    type: STORE_CELL_OUTPUT_SNAPSHOT,
+    payload,
+  };
+};
+
+export const STORE_NOTEBOOK_SNAPSHOT = "STORE_NOTEBOOK_SNAPSHOT";
+export interface StoreNotebookSnapshotAction {
+  type: "STORE_NOTEBOOK_SNAPSHOT";
+  payload: {
+    imageSrc: string;
+    requestId: string;
+  };
+}
+
+export const storeNotebookSnapshot = (payload: {
+  imageSrc: string;
+  requestId: string;
+}): StoreNotebookSnapshotAction => {
+  return {
+    type: STORE_NOTEBOOK_SNAPSHOT,
+    payload,
+  };
+};
+
+export const TAKE_NOTEBOOK_SNAPSHOT = "TAKE_NOTEBOOK_SNAPSHOT";
+export interface TakeNotebookSnapshotAction {
+  type: "TAKE_NOTEBOOK_SNAPSHOT";
+  payload: SnapshotRequest;
+}
+
+export const takeNotebookSnapshot = (payload: SnapshotRequest): TakeNotebookSnapshotAction => {
+  return {
+    type: TAKE_NOTEBOOK_SNAPSHOT,
+    payload,
+  };
+};
+
+export const NOTEBOOK_SNAPSHOT_ERROR = "NOTEBOOK_SNAPSHOT_ERROR";
+export interface NotebookSnapshotErrorAction {
+  type: "NOTEBOOK_SNAPSHOT_ERROR";
+  payload: {
+    error: string;
+  };
+}
+
+export const notebookSnapshotError = (payload: { error: string }): NotebookSnapshotErrorAction => {
+  return {
+    type: NOTEBOOK_SNAPSHOT_ERROR,
     payload,
   };
 };

src/Explorer/Notebook/NotebookComponent/reducers.ts

@@ -70,17 +70,32 @@ export const cdbReducer = (state: CdbRecord, action: Action) => {
       return state.set("hoveredCellId", typedAction.payload.cellId);
     }
 
-    case cdbActions.UPDATE_NOTEBOOK_PARENT_DOM_ELTS: {
-      const typedAction = action as cdbActions.UpdateNotebookParentDomEltAction;
-      var parentEltsMap = state.get("currentNotebookParentElements");
-      const contentRef = typedAction.payload.contentRef;
-      const parentElt = typedAction.payload.parentElt;
-      if (parentElt) {
-        parentEltsMap.set(contentRef, parentElt);
-      } else {
-        parentEltsMap.delete(contentRef);
-      }
-      return state.set("currentNotebookParentElements", parentEltsMap);
+    case cdbActions.STORE_CELL_OUTPUT_SNAPSHOT: {
+      const typedAction = action as cdbActions.StoreCellOutputSnapshotAction;
+      state.cellOutputSnapshots.set(typedAction.payload.cellId, typedAction.payload.snapshot);
+      // TODO Simpler datastructure to instantiate new Map?
+      return state.set("cellOutputSnapshots", new Map(state.cellOutputSnapshots));
+    }
+
+    case cdbActions.STORE_NOTEBOOK_SNAPSHOT: {
+      const typedAction = action as cdbActions.StoreNotebookSnapshotAction;
+      // Clear pending request
+      return state.set("notebookSnapshot", typedAction.payload).set("pendingSnapshotRequest", undefined);
+    }
+
+    case cdbActions.TAKE_NOTEBOOK_SNAPSHOT: {
+      const typedAction = action as cdbActions.TakeNotebookSnapshotAction;
+      // Clear previous snapshots
+      return state
+        .set("cellOutputSnapshots", new Map())
+        .set("notebookSnapshot", undefined)
+        .set("notebookSnapshotError", undefined)
+        .set("pendingSnapshotRequest", typedAction.payload);
+    }
+
+    case cdbActions.NOTEBOOK_SNAPSHOT_ERROR: {
+      const typedAction = action as cdbActions.NotebookSnapshotErrorAction;
+      return state.set("notebookSnapshotError", typedAction.payload.error);
     }
   }
   return state;

src/Explorer/Notebook/NotebookComponent/types.ts

@@ -1,15 +1,40 @@
-import * as Immutable from "immutable";
-import { AppState, ContentRef } from "@nteract/core";
-
-import { Notebook } from "../../../Common/Constants";
 import { CellId } from "@nteract/commutable";
+import { AppState } from "@nteract/core";
+import * as Immutable from "immutable";
+import { Notebook } from "../../../Common/Constants";
+
+export interface SnapshotFragment {
+  image: HTMLImageElement;
+  boundingClientRect: DOMRect;
+  requestId: string;
+}
+
+export type SnapshotRequest = NotebookSnapshotRequest | CellSnapshotRequest;
+interface NotebookSnapshotRequestBase {
+  requestId: string;
+  aspectRatio: number;
+  notebookContentRef: string; // notebook redux contentRef
+  downloadFilename?: string; // Optional: will download as a file
+}
+
+interface NotebookSnapshotRequest extends NotebookSnapshotRequestBase {
+  type: "notebook";
+}
+
+interface CellSnapshotRequest extends NotebookSnapshotRequestBase {
+  type: "celloutput";
+  cellId: string;
+}
 
 export interface CdbRecordProps {
   databaseAccountName: string | undefined;
   defaultExperience: string | undefined;
   kernelRestartDelayMs: number;
   hoveredCellId: CellId | undefined;
-  currentNotebookParentElements: Map<ContentRef, HTMLElement>;
+  cellOutputSnapshots: Map<string, SnapshotFragment>;
+  notebookSnapshot?: { imageSrc: string; requestId: string };
+  pendingSnapshotRequest?: SnapshotRequest;
+  notebookSnapshotError?: string;
 }
 
 export type CdbRecord = Immutable.RecordOf<CdbRecordProps>;
@@ -23,5 +48,8 @@ export const makeCdbRecord = Immutable.Record<CdbRecordProps>({
   defaultExperience: undefined,
   kernelRestartDelayMs: Notebook.kernelRestartInitialDelayMs,
   hoveredCellId: undefined,
-  currentNotebookParentElements: new Map<ContentRef, HTMLElement>(),
+  cellOutputSnapshots: new Map(),
+  notebookSnapshot: undefined,
+  pendingSnapshotRequest: undefined,
+  notebookSnapshotError: undefined,
 });

src/Explorer/Notebook/NotebookManager.tsx

@@ -26,6 +26,7 @@ import { CopyNotebookPane } from "../Panes/CopyNotebookPane/CopyNotebookPane";
 import { PublishNotebookPane } from "../Panes/PublishNotebookPane/PublishNotebookPane";
 import { ResourceTreeAdapter } from "../Tree/ResourceTreeAdapter";
 import { NotebookContentProvider } from "./NotebookComponent/NotebookContentProvider";
+import { SnapshotRequest } from "./NotebookComponent/types";
 import { NotebookContainerClient } from "./NotebookContainerClient";
 import { NotebookContentClient } from "./NotebookContentClient";
 
@@ -112,11 +113,13 @@ export default class NotebookManager {
   public async openPublishNotebookPane(
     name: string,
     content: NotebookPaneContent,
-    parentDomElement: HTMLElement
+    notebookContentRef: string,
+    onTakeSnapshot: (request: SnapshotRequest) => void,
+    onClosePanel: () => void
   ): Promise<void> {
     const explorer = this.params.container;
     explorer.openSidePanel(
-      "New Collection",
+      "Publish Notebook",
       <PublishNotebookPane
         explorer={this.params.container}
         junoClient={this.junoClient}
@@ -125,8 +128,10 @@ export default class NotebookManager {
         name={name}
         author={getFullName()}
         notebookContent={content}
-        parentDomElement={parentDomElement}
-      />
+        notebookContentRef={notebookContentRef}
+        onTakeSnapshot={onTakeSnapshot}
+      />,
+      onClosePanel
     );
   }
 

src/Explorer/Notebook/NotebookRenderer/NotebookRenderer.tsx

@@ -1,6 +1,6 @@
 import { CellId } from "@nteract/commutable";
 import { CellType } from "@nteract/commutable/src";
-import { actions, ContentRef } from "@nteract/core";
+import { actions, ContentRef, selectors } from "@nteract/core";
 import { Cells, CodeCell, RawCell } from "@nteract/stateful-components";
 import MonacoEditor from "@nteract/stateful-components/lib/inputs/connected-editors/monacoEditor";
 import { PassedEditorProps } from "@nteract/stateful-components/lib/inputs/editor";
@@ -12,6 +12,8 @@ import { Dispatch } from "redux";
 import { userContext } from "../../../UserContext";
 import * as cdbActions from "../NotebookComponent/actions";
 import loadTransform from "../NotebookComponent/loadTransform";
+import { CdbAppState, SnapshotFragment, SnapshotRequest } from "../NotebookComponent/types";
+import { NotebookUtil } from "../NotebookUtil";
 import { AzureTheme } from "./AzureTheme";
 import "./base.css";
 import CellCreator from "./decorators/CellCreator";
@@ -32,10 +34,18 @@ export interface NotebookRendererBaseProps {
 }
 
 interface NotebookRendererDispatchProps {
-  updateNotebookParentDomElt: (contentRef: ContentRef, parentElt: HTMLElement) => void;
+  storeNotebookSnapshot: (imageSrc: string, requestId: string) => void;
+  notebookSnapshotError: (error: string) => void;
 }
 
-type NotebookRendererProps = NotebookRendererBaseProps & NotebookRendererDispatchProps;
+interface StateProps {
+  pendingSnapshotRequest: SnapshotRequest;
+  cellOutputSnapshots: Map<string, SnapshotFragment>;
+  notebookSnapshot: { imageSrc: string; requestId: string };
+  nbCodeCells: number;
+}
+
+type NotebookRendererProps = NotebookRendererBaseProps & NotebookRendererDispatchProps & StateProps;
 
 const decorate = (id: string, contentRef: ContentRef, cell_type: CellType, children: React.ReactNode) => {
   const Cell = () => (
@@ -60,27 +70,37 @@ const decorate = (id: string, contentRef: ContentRef, cell_type: CellType, child
 class BaseNotebookRenderer extends React.Component<NotebookRendererProps> {
   private notebookRendererRef = React.createRef<HTMLDivElement>();
 
-  constructor(props: NotebookRendererProps) {
-    super(props);
-
-    this.state = {
-      hoveredCellId: undefined,
-    };
-  }
-
   componentDidMount() {
     if (!userContext.features.sandboxNotebookOutputs) {
       loadTransform(this.props as any);
     }
-    this.props.updateNotebookParentDomElt(this.props.contentRef, this.notebookRendererRef.current);
   }
 
-  componentDidUpdate() {
-    this.props.updateNotebookParentDomElt(this.props.contentRef, this.notebookRendererRef.current);
-  }
-
-  componentWillUnmount() {
-    this.props.updateNotebookParentDomElt(this.props.contentRef, undefined);
+  async componentDidUpdate(): Promise<void> {
+    // Take a snapshot if there's a pending request and all the outputs are also saved
+    if (
+      this.props.pendingSnapshotRequest &&
+      this.props.pendingSnapshotRequest.type === "notebook" &&
+      this.props.pendingSnapshotRequest.notebookContentRef === this.props.contentRef &&
+      (!this.props.notebookSnapshot ||
+        this.props.pendingSnapshotRequest.requestId !== this.props.notebookSnapshot.requestId) &&
+      this.props.cellOutputSnapshots.size === this.props.nbCodeCells
+    ) {
+      try {
+        // Use Html2Canvas because it is much more reliable and fast than dom-to-file
+        const result = await NotebookUtil.takeScreenshotHtml2Canvas(
+          this.notebookRendererRef.current,
+          this.props.pendingSnapshotRequest.aspectRatio,
+          [...this.props.cellOutputSnapshots.values()],
+          this.props.pendingSnapshotRequest.downloadFilename
+        );
+        this.props.storeNotebookSnapshot(result.imageSrc, this.props.pendingSnapshotRequest.requestId);
+      } catch (error) {
+        this.props.notebookSnapshotError(error.message);
+      } finally {
+        this.setState({ processedSnapshotRequest: undefined });
+      }
+    }
   }
 
   render(): JSX.Element {
@@ -156,28 +176,40 @@ class BaseNotebookRenderer extends React.Component<NotebookRendererProps> {
   }
 }
 
+export const makeMapStateToProps = (
+  initialState: CdbAppState,
+  ownProps: NotebookRendererProps
+): ((state: CdbAppState) => StateProps) => {
+  const mapStateToProps = (state: CdbAppState): StateProps => {
+    const { contentRef } = ownProps;
+    const model = selectors.model(state, { contentRef });
+
+    let nbCodeCells;
+    if (model && model.type === "notebook") {
+      nbCodeCells = NotebookUtil.findCodeCellWithDisplay(model.notebook).length;
+    }
+    const { pendingSnapshotRequest, cellOutputSnapshots, notebookSnapshot } = state.cdb;
+    return { pendingSnapshotRequest, cellOutputSnapshots, notebookSnapshot, nbCodeCells };
+  };
+  return mapStateToProps;
+};
+
 const makeMapDispatchToProps = (initialDispatch: Dispatch, initialProps: NotebookRendererBaseProps) => {
   const mapDispatchToProps = (dispatch: Dispatch) => {
     return {
-      addTransform: (transform: React.ComponentType & { MIMETYPE: string }) => {
-        return dispatch(
+      addTransform: (transform: React.ComponentType & { MIMETYPE: string }) =>
+        dispatch(
           actions.addTransform({
             mediaType: transform.MIMETYPE,
             component: transform,
           })
-        );
-      },
-      updateNotebookParentDomElt: (contentRef: ContentRef, parentElt: HTMLElement) => {
-        return dispatch(
-          cdbActions.UpdateNotebookParentDomElt({
-            contentRef,
-            parentElt,
-          })
-        );
-      },
+        ),
+      storeNotebookSnapshot: (imageSrc: string, requestId: string) =>
+        dispatch(cdbActions.storeNotebookSnapshot({ imageSrc, requestId })),
+      notebookSnapshotError: (error: string) => dispatch(cdbActions.notebookSnapshotError({ error })),
     };
   };
   return mapDispatchToProps;
 };
 
-export default connect(null, makeMapDispatchToProps)(BaseNotebookRenderer);
+export default connect(makeMapStateToProps, makeMapDispatchToProps)(BaseNotebookRenderer);

src/Explorer/Notebook/NotebookRenderer/Toolbar.tsx

@@ -1,5 +1,5 @@
 import { ContextualMenuItemType, DirectionalHint, IconButton, IContextualMenuItem } from "@fluentui/react";
-import { CellId, CellType } from "@nteract/commutable";
+import { CellId, CellType, ImmutableCodeCell } from "@nteract/commutable";
 import { actions, AppState, DocumentRecordProps } from "@nteract/core";
 import * as selectors from "@nteract/selectors";
 import { CellToolbarContext } from "@nteract/stateful-components";
@@ -10,6 +10,8 @@ import { connect } from "react-redux";
 import { Dispatch } from "redux";
 import { Action, ActionModifiers } from "../../../Shared/Telemetry/TelemetryConstants";
 import * as cdbActions from "../NotebookComponent/actions";
+import { SnapshotRequest } from "../NotebookComponent/types";
+import { NotebookUtil } from "../NotebookUtil";
 
 export interface ComponentProps {
   contentRef: ContentRef;
@@ -26,12 +28,14 @@ interface DispatchProps {
   clearOutputs: () => void;
   deleteCell: () => void;
   traceNotebookTelemetry: (action: Action, actionModifier?: string, data?: any) => void;
+  takeNotebookSnapshot: (payload: SnapshotRequest) => void;
 }
 
 interface StateProps {
   cellType: CellType;
   cellIdAbove: CellId;
   cellIdBelow: CellId;
+  hasCodeOutput: boolean;
 }
 
 class BaseToolbar extends React.PureComponent<ComponentProps & DispatchProps & StateProps> {
@@ -58,11 +62,29 @@ class BaseToolbar extends React.PureComponent<ComponentProps & DispatchProps & S
             this.props.traceNotebookTelemetry(Action.NotebooksClearOutputsFromMenu, ActionModifiers.Mark);
           },
         },
-        {
-          key: "Divider",
-          itemType: ContextualMenuItemType.Divider,
-        },
       ]);
+
+      if (this.props.hasCodeOutput) {
+        items.push({
+          key: "Export output to image",
+          text: "Export output to image",
+          onClick: () => {
+            this.props.takeNotebookSnapshot({
+              requestId: new Date().getTime().toString(),
+              aspectRatio: undefined,
+              type: "celloutput",
+              cellId: this.props.id,
+              notebookContentRef: this.props.contentRef,
+              downloadFilename: `celloutput-${this.props.contentRef}_${this.props.id}.png`,
+            });
+          },
+        });
+      }
+
+      items.push({
+        key: "Divider",
+        itemType: ContextualMenuItemType.Divider,
+      });
     }
 
     items = items.concat([
@@ -183,12 +205,13 @@ const mapDispatchToProps = (
   deleteCell: () => dispatch(actions.deleteCell({ id, contentRef })),
   traceNotebookTelemetry: (action: Action, actionModifier?: string, data?: any) =>
     dispatch(cdbActions.traceNotebookTelemetry({ action, actionModifier, data })),
+  takeNotebookSnapshot: (request: SnapshotRequest) => dispatch(cdbActions.takeNotebookSnapshot(request)),
 });
 
 const makeMapStateToProps = (state: AppState, ownProps: ComponentProps): ((state: AppState) => StateProps) => {
   const mapStateToProps = (state: AppState) => {
-    const cellType = selectors.cell.cellFromState(state, { id: ownProps.id, contentRef: ownProps.contentRef })
-      .cell_type;
+    const cell = selectors.cell.cellFromState(state, { id: ownProps.id, contentRef: ownProps.contentRef });
+    const cellType = cell.cell_type;
     const model = selectors.model(state, { contentRef: ownProps.contentRef });
     const cellOrder = selectors.notebook.cellOrder(model as RecordOf<DocumentRecordProps>);
     const cellIndex = cellOrder.indexOf(ownProps.id);
@@ -199,6 +222,7 @@ const makeMapStateToProps = (state: AppState, ownProps: ComponentProps): ((state
       cellType,
       cellIdAbove,
       cellIdBelow,
+      hasCodeOutput: cellType === "code" && NotebookUtil.hasCodeCellOutput(cell as ImmutableCodeCell),
     };
   };
   return mapStateToProps;

src/Explorer/Notebook/NotebookRenderer/outputs/SandboxOutputs.tsx

@@ -7,7 +7,9 @@ import postRobot from "post-robot";
 import React from "react";
 import { connect } from "react-redux";
 import { Dispatch } from "redux";
-import { CellOutputViewerProps } from "../../../../CellOutputViewer/CellOutputViewer";
+import { CellOutputViewerProps, SnapshotResponse } from "../../../../CellOutputViewer/CellOutputViewer";
+import * as cdbActions from "../../NotebookComponent/actions";
+import { CdbAppState, SnapshotFragment, SnapshotRequest } from "../../NotebookComponent/types";
 
 // Adapted from https://github.com/nteract/nteract/blob/main/packages/stateful-components/src/outputs/index.tsx
 // to add support for sandboxing using <iframe>
@@ -24,27 +26,47 @@ interface StateProps {
   expanded: boolean;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   outputs: Immutable.List<any>;
+
+  pendingSnapshotRequest: SnapshotRequest;
 }
 
 interface DispatchProps {
   onMetadataChange?: (metadata: JSONObject, mediaType: string, index?: number) => void;
+  storeNotebookSnapshot: (imageSrc: string, requestId: string) => void;
+  storeSnapshotFragment: (cellId: string, snapshotFragment: SnapshotFragment) => void;
+  notebookSnapshotError: (error: string) => void;
 }
 
-export class SandboxOutputs extends React.PureComponent<ComponentProps & StateProps & DispatchProps> {
+type SandboxOutputsProps = ComponentProps & StateProps & DispatchProps;
+
+export class SandboxOutputs extends React.Component<SandboxOutputsProps> {
   private childWindow: Window;
+  private nodeRef = React.createRef<HTMLDivElement>();
+
+  constructor(props: SandboxOutputsProps) {
+    super(props);
+
+    this.state = {
+      processedSnapshotRequest: undefined,
+    };
+  }
 
   render(): JSX.Element {
     // Using min-width to set the width of the iFrame, works around an issue in iOS that can prevent the iFrame from sizing correctly.
-    return (
-      <IframeResizer
-        checkOrigin={false}
-        loading="lazy"
-        heightCalculationMethod="taggedElement"
-        onLoad={(event) => this.handleFrameLoad(event)}
-        src="./cellOutputViewer.html"
-        style={{ height: "1px", width: "1px", minWidth: "100%", border: "none" }}
-        sandbox="allow-downloads allow-popups allow-forms allow-pointer-lock allow-scripts allow-popups-to-escape-sandbox"
-      />
+    return this.props.outputs && this.props.outputs.size > 0 ? (
+      <div ref={this.nodeRef}>
+        <IframeResizer
+          checkOrigin={false}
+          loading="lazy"
+          heightCalculationMethod="taggedElement"
+          onLoad={(event) => this.handleFrameLoad(event)}
+          src="./cellOutputViewer.html"
+          style={{ height: "1px", width: "1px", minWidth: "100%", border: "none" }}
+          sandbox="allow-downloads allow-popups allow-forms allow-pointer-lock allow-scripts allow-popups-to-escape-sandbox"
+        />
+      </div>
+    ) : (
+      <></>
     );
   }
 
@@ -76,8 +98,48 @@ export class SandboxOutputs extends React.PureComponent<ComponentProps & StatePr
     this.sendPropsToFrame();
   }
 
-  componentDidUpdate(): void {
+  async componentDidUpdate(prevProps: SandboxOutputsProps): Promise<void> {
     this.sendPropsToFrame();
+
+    if (
+      this.props.pendingSnapshotRequest &&
+      prevProps.pendingSnapshotRequest !== this.props.pendingSnapshotRequest &&
+      this.props.pendingSnapshotRequest.notebookContentRef === this.props.contentRef &&
+      this.nodeRef?.current
+    ) {
+      const boundingClientRect = this.nodeRef.current.getBoundingClientRect();
+
+      try {
+        const { data } = (await postRobot.send(
+          this.childWindow,
+          "snapshotRequest",
+          this.props.pendingSnapshotRequest
+        )) as { data: SnapshotResponse };
+        if (this.props.pendingSnapshotRequest.type === "notebook") {
+          if (data.imageSrc === undefined) {
+            this.props.storeSnapshotFragment(this.props.id, {
+              image: undefined,
+              boundingClientRect: boundingClientRect,
+              requestId: data.requestId,
+            });
+            return;
+          }
+          const image = new Image();
+          image.src = data.imageSrc;
+          image.onload = () => {
+            this.props.storeSnapshotFragment(this.props.id, {
+              image,
+              boundingClientRect: boundingClientRect,
+              requestId: data.requestId,
+            });
+          };
+        } else if (this.props.pendingSnapshotRequest.type === "celloutput") {
+          this.props.storeNotebookSnapshot(data.imageSrc, this.props.pendingSnapshotRequest.requestId);
+        }
+      } catch (error) {
+        this.props.notebookSnapshotError(error.message);
+      }
+    }
   }
 }
 
@@ -85,7 +147,7 @@ export const makeMapStateToProps = (
   initialState: AppState,
   ownProps: ComponentProps
 ): ((state: AppState) => StateProps) => {
-  const mapStateToProps = (state: AppState): StateProps => {
+  const mapStateToProps = (state: CdbAppState): StateProps => {
     let outputs = Immutable.List();
     let hidden = false;
     let expanded = false;
@@ -102,7 +164,17 @@ export const makeMapStateToProps = (
       }
     }
 
-    return { outputs, hidden, expanded };
+    // Determine whether to take a snapshot or not
+    let pendingSnapshotRequest = state.cdb.pendingSnapshotRequest;
+    if (
+      pendingSnapshotRequest &&
+      pendingSnapshotRequest.type === "celloutput" &&
+      pendingSnapshotRequest.cellId !== id
+    ) {
+      pendingSnapshotRequest = undefined;
+    }
+
+    return { outputs, hidden, expanded, pendingSnapshotRequest };
   };
   return mapStateToProps;
 };
@@ -125,6 +197,11 @@ export const makeMapDispatchToProps = (
           })
         );
       },
+      storeSnapshotFragment: (cellId: string, snapshot: SnapshotFragment) =>
+        dispatch(cdbActions.storeCellOutputSnapshot({ cellId, snapshot })),
+      storeNotebookSnapshot: (imageSrc: string, requestId: string) =>
+        dispatch(cdbActions.storeNotebookSnapshot({ imageSrc, requestId })),
+      notebookSnapshotError: (error: string) => dispatch(cdbActions.notebookSnapshotError({ error })),
     };
   };
   return mapDispatchToProps;

src/Explorer/Notebook/NotebookUtil.test.ts

@@ -1,15 +1,15 @@
-import { NotebookUtil } from "./NotebookUtil";
-import * as GitHubUtils from "../../Utils/GitHubUtils";
 import {
-  ImmutableNotebook,
-  MediaBundle,
   CodeCellParams,
-  MarkdownCellParams,
+  ImmutableNotebook,
   makeCodeCell,
   makeMarkdownCell,
   makeNotebookRecord,
+  MarkdownCellParams,
+  MediaBundle,
 } from "@nteract/commutable";
 import { List, Map } from "immutable";
+import * as GitHubUtils from "../../Utils/GitHubUtils";
+import { NotebookUtil } from "./NotebookUtil";
 
 const fileName = "file";
 const notebookName = "file.ipynb";
@@ -131,7 +131,7 @@ describe("NotebookUtil", () => {
   describe("findFirstCodeCellWithDisplay", () => {
     it("works for Notebook file", () => {
       const notebookObject = notebookRecord as ImmutableNotebook;
-      expect(NotebookUtil.findFirstCodeCellWithDisplay(notebookObject)).toEqual(1);
+      expect(NotebookUtil.findCodeCellWithDisplay(notebookObject)[0]).toEqual("1");
     });
   });
 });

src/Explorer/Notebook/NotebookUtil.ts

@@ -1,8 +1,11 @@
+import { ImmutableCodeCell, ImmutableNotebook } from "@nteract/commutable";
+import domtoimage from "dom-to-image";
+import Html2Canvas from "html2canvas";
 import path from "path";
-import { ImmutableNotebook, ImmutableCodeCell } from "@nteract/commutable";
-import { NotebookContentItem, NotebookContentItemType } from "./NotebookContentItem";
-import * as StringUtils from "../../Utils/StringUtils";
 import * as GitHubUtils from "../../Utils/GitHubUtils";
+import * as StringUtils from "../../Utils/StringUtils";
+import { SnapshotFragment } from "./NotebookComponent/types";
+import { NotebookContentItem, NotebookContentItemType } from "./NotebookContentItem";
 
 // Must match rx-jupyter' FileType
 export type FileType = "directory" | "file" | "notebook";
@@ -141,23 +144,175 @@ export class NotebookUtil {
     return `${basePath}${newName}`;
   }
 
-  public static findFirstCodeCellWithDisplay(notebookObject: ImmutableNotebook): number {
-    let codeCellIndex = 0;
-    for (let i = 0; i < notebookObject.cellOrder.size; i++) {
-      const cellId = notebookObject.cellOrder.get(i);
-      if (cellId) {
-        const cell = notebookObject.cellMap.get(cellId);
-        if (cell?.cell_type === "code") {
-          const displayOutput = (cell as ImmutableCodeCell)?.outputs?.find(
-            (output) => output.output_type === "display_data" || output.output_type === "execute_result"
-          );
-          if (displayOutput) {
-            return codeCellIndex;
-          }
-          codeCellIndex++;
+  public static hasCodeCellOutput(cell: ImmutableCodeCell): boolean {
+    return !!cell?.outputs?.find(
+      (output) =>
+        output.output_type === "display_data" ||
+        output.output_type === "execute_result" ||
+        output.output_type === "stream"
+    );
+  }
+
+  /**
+   * Find code cells with display
+   * @param notebookObject
+   * @returns array of cell ids
+   */
+  public static findCodeCellWithDisplay(notebookObject: ImmutableNotebook): string[] {
+    return notebookObject.cellOrder.reduce((accumulator: string[], cellId) => {
+      const cell = notebookObject.cellMap.get(cellId);
+      if (cell?.cell_type === "code") {
+        if (NotebookUtil.hasCodeCellOutput(cell as ImmutableCodeCell)) {
+          accumulator.push(cellId);
         }
       }
-    }
-    throw new Error("Output does not exist for any of the cells.");
+      return accumulator;
+    }, []);
+  }
+
+  public static takeScreenshotHtml2Canvas = (
+    target: HTMLElement,
+    aspectRatio: number,
+    subSnapshots: SnapshotFragment[],
+    downloadFilename?: string
+  ): Promise<{ imageSrc: string | undefined }> => {
+    return new Promise(async (resolve, reject) => {
+      try {
+        // target.scrollIntoView();
+        const canvas = await Html2Canvas(target, {
+          useCORS: true,
+          allowTaint: true,
+          scale: 1,
+          logging: false,
+        });
+
+        //redraw canvas to fit aspect ratio
+        const originalImageData = canvas.toDataURL();
+        const width = parseInt(canvas.style.width.split("px")[0]);
+        if (aspectRatio) {
+          canvas.height = width * aspectRatio;
+        }
+
+        if (originalImageData === "data:,") {
+          // Empty output
+          resolve({ imageSrc: undefined });
+          return;
+        }
+
+        const context = canvas.getContext("2d");
+        const image = new Image();
+        image.src = originalImageData;
+        image.onload = () => {
+          if (!context) {
+            reject(new Error("No context to draw on"));
+            return;
+          }
+          context.drawImage(image, 0, 0);
+
+          // draw sub images
+          if (subSnapshots) {
+            const parentRect = target.getBoundingClientRect();
+            subSnapshots.forEach((snapshot) => {
+              if (snapshot.image) {
+                context.drawImage(
+                  snapshot.image,
+                  snapshot.boundingClientRect.x - parentRect.x,
+                  snapshot.boundingClientRect.y - parentRect.y
+                );
+              }
+            });
+          }
+
+          resolve({ imageSrc: canvas.toDataURL() });
+
+          if (downloadFilename) {
+            NotebookUtil.downloadFile(
+              downloadFilename,
+              canvas.toDataURL("image/png").replace("image/png", "image/octet-stream")
+            );
+          }
+        };
+      } catch (error) {
+        return reject(error);
+      }
+    });
+  };
+
+  public static takeScreenshotDomToImage = (
+    target: HTMLElement,
+    aspectRatio: number,
+    subSnapshots: SnapshotFragment[],
+    downloadFilename?: string
+  ): Promise<{ imageSrc?: string }> => {
+    return new Promise(async (resolve, reject) => {
+      // target.scrollIntoView();
+      try {
+        const filter = (node: Node): boolean => {
+          const excludedList = ["IMG", "CANVAS"];
+          return !excludedList.includes((node as HTMLElement).tagName);
+        };
+
+        const originalImageData = await domtoimage.toPng(target, { filter });
+        if (originalImageData === "data:,") {
+          // Empty output
+          resolve({});
+          return;
+        }
+
+        const baseImage = new Image();
+        baseImage.src = originalImageData;
+        baseImage.onload = () => {
+          const canvas = document.createElement("canvas");
+          canvas.width = baseImage.width;
+          canvas.height = aspectRatio !== undefined ? baseImage.width * aspectRatio : baseImage.width;
+
+          const context = canvas.getContext("2d");
+          if (!context) {
+            reject(new Error("No Canvas to draw on"));
+            return;
+          }
+
+          // White background otherwise image is transparent
+          context.fillStyle = "white";
+          context.fillRect(0, 0, baseImage.width, baseImage.height);
+
+          context.drawImage(baseImage, 0, 0);
+
+          // draw sub images
+          if (subSnapshots) {
+            const parentRect = target.getBoundingClientRect();
+            subSnapshots.forEach((snapshot) => {
+              if (snapshot.image) {
+                context.drawImage(
+                  snapshot.image,
+                  snapshot.boundingClientRect.x - parentRect.x,
+                  snapshot.boundingClientRect.y - parentRect.y
+                );
+              }
+            });
+          }
+
+          resolve({ imageSrc: canvas.toDataURL() });
+
+          if (downloadFilename) {
+            NotebookUtil.downloadFile(
+              downloadFilename,
+              canvas.toDataURL("image/png").replace("image/png", "image/octet-stream")
+            );
+          }
+        };
+      } catch (error) {
+        reject(error);
+      }
+    });
+  };
+
+  private static downloadFile(filename: string, content: string): void {
+    const link = document.createElement("a");
+    link.href = content;
+    link.download = filename;
+    document.body.appendChild(link);
+    link.click();
+    document.body.removeChild(link);
   }
 }