From fd81d2a22dd54c483414e5466b60ba53a14439ad Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Mon, 30 May 2022 16:40:19 +1000
Subject: [PATCH 1/5] add legend to zoomed in plot

---
 webview/src/plots/components/ZoomablePlot.tsx | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 1bd6f73d69..58e3747a77 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -31,7 +31,17 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
     }
   }, [plotProps, id, renderZoomedInPlot])
 
-  const handleOnClick = () => renderZoomedInPlot(plotProps, id)
+  const handleOnClick = () => {
+    const zoomedPlotWithLegend = { ...plotProps } as VegaLiteProps & {
+      spec: { encoding?: { color?: { legend?: { disable?: boolean } } } }
+    }
+    if (
+      zoomedPlotWithLegend.spec.encoding?.color?.legend?.disable !== undefined
+    ) {
+      delete zoomedPlotWithLegend.spec.encoding.color.legend.disable
+    }
+    return renderZoomedInPlot(zoomedPlotWithLegend, id)
+  }
 
   return (
     <button className={styles.zoomablePlot} onClick={handleOnClick}>

From 170100dc0a27ca4fb9f63f86cab8adf31968542f Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 31 May 2022 16:25:32 +1000
Subject: [PATCH 2/5] add legend for zoomed in checkpoint plots

---
 webview/src/plots/components/ZoomablePlot.tsx |  25 +--
 .../plots/components/checkpointPlots/util.ts  | 150 +++++++++---------
 2 files changed, 93 insertions(+), 82 deletions(-)

diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 58e3747a77..d8700d63cb 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -16,6 +16,20 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
+const removeLegendSuppression = (plotProps: VegaLiteProps): VegaLiteProps => {
+  const plot = { ...plotProps } as VegaLiteProps & {
+    spec: {
+      encoding?: {
+        color?: { legend?: { disable?: boolean } }
+      }
+    }
+  }
+  if (plot.spec.encoding?.color?.legend?.disable !== undefined) {
+    delete plot.spec.encoding.color.legend.disable
+  }
+  return plot
+}
+
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
@@ -32,15 +46,8 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   }, [plotProps, id, renderZoomedInPlot])
 
   const handleOnClick = () => {
-    const zoomedPlotWithLegend = { ...plotProps } as VegaLiteProps & {
-      spec: { encoding?: { color?: { legend?: { disable?: boolean } } } }
-    }
-    if (
-      zoomedPlotWithLegend.spec.encoding?.color?.legend?.disable !== undefined
-    ) {
-      delete zoomedPlotWithLegend.spec.encoding.color.legend.disable
-    }
-    return renderZoomedInPlot(zoomedPlotWithLegend, id)
+    const zoomedPlot = removeLegendSuppression(plotProps)
+    return renderZoomedInPlot(zoomedPlot, id)
   }
 
   return (
diff --git a/webview/src/plots/components/checkpointPlots/util.ts b/webview/src/plots/components/checkpointPlots/util.ts
index 2c6ad57def..94e099555b 100644
--- a/webview/src/plots/components/checkpointPlots/util.ts
+++ b/webview/src/plots/components/checkpointPlots/util.ts
@@ -4,86 +4,90 @@ import { ColorScale } from 'dvc/src/plots/webview/contract'
 export const createSpec = (
   title: string,
   scale?: ColorScale
-): VisualizationSpec => ({
-  $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
-  data: { name: 'values' },
-  encoding: {
-    x: {
-      axis: { format: '0d', tickMinStep: 1 },
-      field: 'iteration',
-      title: 'iteration',
-      type: 'quantitative'
+): VisualizationSpec =>
+  ({
+    $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
+    data: { name: 'values' },
+    encoding: {
+      color: {
+        field: 'group',
+        legend: { disable: true },
+        scale,
+        title: 'rev',
+        type: 'nominal'
+      },
+      x: {
+        axis: { format: '0d', tickMinStep: 1 },
+        field: 'iteration',
+        title: 'iteration',
+        type: 'quantitative'
+      },
+      y: {
+        field: 'y',
+        scale: { zero: false },
+        title,
+        type: 'quantitative'
+      }
     },
-    y: {
-      field: 'y',
-      scale: { zero: false },
-      title,
-      type: 'quantitative'
-    }
-  },
-  height: 'container',
-  layer: [
-    {
-      encoding: {
-        color: { field: 'group', legend: null, scale, type: 'nominal' }
+    height: 'container',
+    layer: [
+      {
+        layer: [
+          { mark: { type: 'line' } },
+          {
+            mark: { type: 'point' },
+            transform: [
+              {
+                filter: { empty: false, param: 'hover' }
+              }
+            ]
+          }
+        ]
       },
-
-      layer: [
-        { mark: { type: 'line' } },
-        {
-          mark: { type: 'point' },
-          transform: [
+      {
+        encoding: {
+          opacity: { value: 0 },
+          tooltip: [
+            { field: 'group', title: 'name' },
             {
-              filter: { empty: false, param: 'hover' }
+              field: 'y',
+              title: title.slice(Math.max(0, title.indexOf(':') + 1)),
+              type: 'quantitative'
             }
           ]
-        }
-      ]
-    },
-    {
-      encoding: {
-        opacity: { value: 0 },
-        tooltip: [
-          { field: 'group', title: 'name' },
+        },
+        mark: { type: 'rule' },
+        params: [
           {
-            field: 'y',
-            title: title.slice(Math.max(0, title.indexOf(':') + 1)),
-            type: 'quantitative'
+            name: 'hover',
+            select: {
+              clear: 'mouseout',
+              fields: ['iteration', 'y'],
+              nearest: true,
+              on: 'mouseover',
+              type: 'point'
+            }
           }
         ]
       },
-      mark: { type: 'rule' },
-      params: [
-        {
-          name: 'hover',
-          select: {
-            clear: 'mouseout',
-            fields: ['iteration', 'y'],
-            nearest: true,
-            on: 'mouseover',
-            type: 'point'
+      {
+        encoding: {
+          color: { field: 'group', scale },
+          x: { aggregate: 'max', field: 'iteration', type: 'quantitative' },
+          y: {
+            aggregate: { argmax: 'iteration' },
+            field: 'y',
+            type: 'quantitative'
           }
-        }
-      ]
-    },
-    {
-      encoding: {
-        color: { field: 'group', scale },
-        x: { aggregate: 'max', field: 'iteration', type: 'quantitative' },
-        y: {
-          aggregate: { argmax: 'iteration' },
-          field: 'y',
-          type: 'quantitative'
-        }
-      },
-      mark: { stroke: null, type: 'circle' }
-    }
-  ],
-  transform: [
-    {
-      as: 'y',
-      calculate: "format(datum['y'],'.5f')"
-    }
-  ],
-  width: 'container'
-})
+        },
+        mark: { stroke: null, type: 'circle' }
+      }
+    ],
+    transform: [
+      {
+        as: 'y',
+        calculate: "format(datum['y'],'.5f')"
+      }
+    ],
+    width: 'container'
+  } as VisualizationSpec)

From adf92cf34eefb2b816acb3ee4f7bba504e6113c8 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 31 May 2022 21:54:45 +1000
Subject: [PATCH 3/5] refactor removal of legend suppression up to top level
 function

---
 webview/src/plots/components/Plots.tsx        |  8 +++++++-
 webview/src/plots/components/ZoomablePlot.tsx | 19 +------------------
 webview/src/plots/components/util.ts          | 17 +++++++++++++++++
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/webview/src/plots/components/Plots.tsx b/webview/src/plots/components/Plots.tsx
index db96be8223..fc333d3658 100644
--- a/webview/src/plots/components/Plots.tsx
+++ b/webview/src/plots/components/Plots.tsx
@@ -4,6 +4,7 @@ import React, { useEffect, useRef, useState, useCallback } from 'react'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import { Config } from 'vega-lite'
 import styles from './styles.module.scss'
+import { removeLegendSuppression } from './util'
 import { PlotsSizeProvider } from './PlotsSizeContext'
 import { AddPlots, Welcome } from './GetStarted'
 import { CheckpointPlotsWrapper } from './checkpointPlots/CheckpointPlotsWrapper'
@@ -40,10 +41,15 @@ const PlotsContent = ({ state }: PlotsProps) => {
     }
   }, [zoomedInPlot])
 
+  const setZoomedInPlotWithLegend = (props?: VegaLiteProps) => {
+    removeLegendSuppression(props)
+    return setZoomedInPlot(props)
+  }
+
   const handleZoomInPlot = useCallback(
     (props: VegaLiteProps, id: string, refresh?: boolean) => {
       if (!refresh) {
-        setZoomedInPlot(props)
+        setZoomedInPlotWithLegend(props)
         zoomedInPlotId.current = id
         return
       }
diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index d8700d63cb..1bd6f73d69 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -16,20 +16,6 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
-const removeLegendSuppression = (plotProps: VegaLiteProps): VegaLiteProps => {
-  const plot = { ...plotProps } as VegaLiteProps & {
-    spec: {
-      encoding?: {
-        color?: { legend?: { disable?: boolean } }
-      }
-    }
-  }
-  if (plot.spec.encoding?.color?.legend?.disable !== undefined) {
-    delete plot.spec.encoding.color.legend.disable
-  }
-  return plot
-}
-
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
@@ -45,10 +31,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
     }
   }, [plotProps, id, renderZoomedInPlot])
 
-  const handleOnClick = () => {
-    const zoomedPlot = removeLegendSuppression(plotProps)
-    return renderZoomedInPlot(zoomedPlot, id)
-  }
+  const handleOnClick = () => renderZoomedInPlot(plotProps, id)
 
   return (
     <button className={styles.zoomablePlot} onClick={handleOnClick}>
diff --git a/webview/src/plots/components/util.ts b/webview/src/plots/components/util.ts
index 96bf5e46bf..768e5e6f8b 100644
--- a/webview/src/plots/components/util.ts
+++ b/webview/src/plots/components/util.ts
@@ -1,4 +1,5 @@
 import { PlotSize } from 'dvc/src/plots/webview/contract'
+import { VegaLiteProps } from 'react-vega/lib/VegaLite'
 
 const MaxItemsBeforeVirtualization = {
   [PlotSize.LARGE]: 10,
@@ -44,3 +45,19 @@ export const getNbItemsPerRow = (size: PlotSize) => {
 
 export const shouldUseVirtualizedGrid = (nbItems: number, size: PlotSize) =>
   nbItems > MaxItemsBeforeVirtualization[size]
+
+export const removeLegendSuppression = (
+  plotProps?: VegaLiteProps
+): VegaLiteProps => {
+  const plot = plotProps as VegaLiteProps & {
+    spec: {
+      encoding?: {
+        color?: { legend?: { disable?: boolean } }
+      }
+    }
+  }
+  if (plot?.spec?.encoding?.color?.legend?.disable !== undefined) {
+    delete plot.spec.encoding.color.legend.disable
+  }
+  return plot
+}

From c78dbbdb26b7ba2aaedc30947bd75394fecf22d6 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Tue, 31 May 2022 23:41:59 +1000
Subject: [PATCH 4/5] commit broken storybook for now

---
 webview/src/plots/components/Plots.tsx        | 11 ++++-------
 webview/src/plots/components/ZoomablePlot.tsx | 12 +++++++++---
 webview/src/plots/components/util.ts          |  4 ++--
 webview/src/stories/Plots.stories.tsx         | 13 +++++++++++++
 4 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/webview/src/plots/components/Plots.tsx b/webview/src/plots/components/Plots.tsx
index fc333d3658..d3cd06db25 100644
--- a/webview/src/plots/components/Plots.tsx
+++ b/webview/src/plots/components/Plots.tsx
@@ -41,21 +41,18 @@ const PlotsContent = ({ state }: PlotsProps) => {
     }
   }, [zoomedInPlot])
 
-  const setZoomedInPlotWithLegend = (props?: VegaLiteProps) => {
-    removeLegendSuppression(props)
-    return setZoomedInPlot(props)
-  }
-
   const handleZoomInPlot = useCallback(
     (props: VegaLiteProps, id: string, refresh?: boolean) => {
       if (!refresh) {
-        setZoomedInPlotWithLegend(props)
+        setZoomedInPlot(removeLegendSuppression(props))
         zoomedInPlotId.current = id
         return
       }
 
       if (zoomedInPlotId.current === id) {
-        setZoomedInPlot(plot => (plot ? props : undefined))
+        setZoomedInPlot(plot =>
+          plot ? removeLegendSuppression(props) : undefined
+        )
       }
     },
     [setZoomedInPlot]
diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index 1bd6f73d69..d5dd7707c5 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -1,6 +1,7 @@
 import React, { useEffect, useRef } from 'react'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import styles from './styles.module.scss'
+import { removeLegendSuppression } from './util'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
 
 export interface ZoomablePlotProps {
@@ -16,6 +17,13 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
+const areDifferentPlotProps = (
+  currentPlotProps: VegaLiteProps,
+  plotProps: VegaLiteProps
+) =>
+  JSON.stringify(removeLegendSuppression(currentPlotProps)) !==
+  JSON.stringify(removeLegendSuppression(plotProps))
+
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
@@ -23,9 +31,7 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
 }) => {
   const previousPlotProps = useRef(plotProps)
   useEffect(() => {
-    if (
-      JSON.stringify(previousPlotProps.current) !== JSON.stringify(plotProps)
-    ) {
+    if (areDifferentPlotProps(previousPlotProps.current, plotProps)) {
       renderZoomedInPlot(plotProps, id, true)
       previousPlotProps.current = plotProps
     }
diff --git a/webview/src/plots/components/util.ts b/webview/src/plots/components/util.ts
index 768e5e6f8b..d9055c86fa 100644
--- a/webview/src/plots/components/util.ts
+++ b/webview/src/plots/components/util.ts
@@ -49,7 +49,7 @@ export const shouldUseVirtualizedGrid = (nbItems: number, size: PlotSize) =>
 export const removeLegendSuppression = (
   plotProps?: VegaLiteProps
 ): VegaLiteProps => {
-  const plot = plotProps as VegaLiteProps & {
+  const plot = { ...plotProps } as VegaLiteProps & {
     spec: {
       encoding?: {
         color?: { legend?: { disable?: boolean } }
@@ -57,7 +57,7 @@ export const removeLegendSuppression = (
     }
   }
   if (plot?.spec?.encoding?.color?.legend?.disable !== undefined) {
-    delete plot.spec.encoding.color.legend.disable
+    plot.spec.encoding.color.legend.disable = false
   }
   return plot
 }
diff --git a/webview/src/stories/Plots.stories.tsx b/webview/src/stories/Plots.stories.tsx
index 85f3db86c0..51eba902bb 100644
--- a/webview/src/stories/Plots.stories.tsx
+++ b/webview/src/stories/Plots.stories.tsx
@@ -185,3 +185,16 @@ MultiviewZoomedInPlot.play = async ({ canvasElement }) => {
 
   fireEvent.click(plotButton)
 }
+
+export const CheckpointZoomedInPlot = Template.bind({})
+CheckpointZoomedInPlot.parameters = {
+  chromatic: { delay: 500 }
+}
+CheckpointZoomedInPlot.play = async ({ canvasElement }) => {
+  const canvas = within(canvasElement)
+  const plot = await canvas.findByText('summary.json:val_accuracy')
+
+  plot.scrollIntoView()
+
+  fireEvent.click(plot)
+}

From 2e404cd0c6903bee320431ccd0f838877647524c Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Wed, 1 Jun 2022 09:58:23 +1000
Subject: [PATCH 5/5] create zoomed in plot component and do not mutate props

---
 webview/package.json                          |  2 ++
 webview/src/plots/components/Plots.tsx        | 29 +++------------
 webview/src/plots/components/ZoomablePlot.tsx | 12 ++-----
 webview/src/plots/components/ZoomedInPlot.tsx | 36 +++++++++++++++++++
 webview/src/plots/components/util.ts          | 17 ---------
 5 files changed, 46 insertions(+), 50 deletions(-)
 create mode 100644 webview/src/plots/components/ZoomedInPlot.tsx

diff --git a/webview/package.json b/webview/package.json
index 95d422d1f2..12c21c3aec 100644
--- a/webview/package.json
+++ b/webview/package.json
@@ -25,6 +25,8 @@
     "@tippyjs/react": "^4.2.6",
     "@vscode/webview-ui-toolkit": "^1.0.0",
     "classnames": "^2.2.6",
+    "lodash.clonedeep": "^4.5.0",
+    "lodash.merge": "^4.6.2",
     "react": "^17.0.1",
     "react-dom": "^17.0.1",
     "react-table": "^7.7.0",
diff --git a/webview/src/plots/components/Plots.tsx b/webview/src/plots/components/Plots.tsx
index d3cd06db25..417f7bd754 100644
--- a/webview/src/plots/components/Plots.tsx
+++ b/webview/src/plots/components/Plots.tsx
@@ -1,12 +1,10 @@
 import { PlotSize, Section } from 'dvc/src/plots/webview/contract'
 import { MessageFromWebviewType } from 'dvc/src/webview/contract'
 import React, { useEffect, useRef, useState, useCallback } from 'react'
-import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
-import { Config } from 'vega-lite'
-import styles from './styles.module.scss'
-import { removeLegendSuppression } from './util'
+import { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import { PlotsSizeProvider } from './PlotsSizeContext'
 import { AddPlots, Welcome } from './GetStarted'
+import { ZoomedInPlot } from './ZoomedInPlot'
 import { CheckpointPlotsWrapper } from './checkpointPlots/CheckpointPlotsWrapper'
 import { TemplatePlotsWrapper } from './templatePlots/TemplatePlotsWrapper'
 import { ComparisonTableWrapper } from './comparisonTable/ComparisonTableWrapper'
@@ -16,7 +14,6 @@ import { Modal } from '../../shared/components/modal/Modal'
 import { WebviewWrapper } from '../../shared/components/webviewWrapper/WebviewWrapper'
 import { DragDropProvider } from '../../shared/components/dragDrop/DragDropContext'
 import { sendMessage } from '../../shared/vscode'
-import { getThemeValue, ThemeProperty } from '../../util/styles'
 import { GetStarted } from '../../shared/components/getStarted/GetStarted'
 
 interface PlotsProps {
@@ -44,15 +41,13 @@ const PlotsContent = ({ state }: PlotsProps) => {
   const handleZoomInPlot = useCallback(
     (props: VegaLiteProps, id: string, refresh?: boolean) => {
       if (!refresh) {
-        setZoomedInPlot(removeLegendSuppression(props))
+        setZoomedInPlot(props)
         zoomedInPlotId.current = id
         return
       }
 
       if (zoomedInPlotId.current === id) {
-        setZoomedInPlot(plot =>
-          plot ? removeLegendSuppression(props) : undefined
-        )
+        setZoomedInPlot(plot => (plot ? props : undefined))
       }
     },
     [setZoomedInPlot]
@@ -152,21 +147,7 @@ const PlotsContent = ({ state }: PlotsProps) => {
 
       {zoomedInPlot && (
         <Modal onClose={handleModalClose}>
-          <div className={styles.zoomedInPlot} data-testid="zoomed-in-plot">
-            <VegaLite
-              {...zoomedInPlot}
-              config={{
-                ...(zoomedInPlot.config as Config),
-                background: getThemeValue(ThemeProperty.MENU_BACKGROUND)
-              }}
-              actions={{
-                compiled: false,
-                editor: false,
-                export: true,
-                source: false
-              }}
-            />
-          </div>
+          <ZoomedInPlot props={zoomedInPlot} />
         </Modal>
       )}
     </DragDropProvider>
diff --git a/webview/src/plots/components/ZoomablePlot.tsx b/webview/src/plots/components/ZoomablePlot.tsx
index d5dd7707c5..1bd6f73d69 100644
--- a/webview/src/plots/components/ZoomablePlot.tsx
+++ b/webview/src/plots/components/ZoomablePlot.tsx
@@ -1,7 +1,6 @@
 import React, { useEffect, useRef } from 'react'
 import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
 import styles from './styles.module.scss'
-import { removeLegendSuppression } from './util'
 import { GripIcon } from '../../shared/components/dragDrop/GripIcon'
 
 export interface ZoomablePlotProps {
@@ -17,13 +16,6 @@ interface ZoomablePlotOwnProps extends ZoomablePlotProps {
   id: string
 }
 
-const areDifferentPlotProps = (
-  currentPlotProps: VegaLiteProps,
-  plotProps: VegaLiteProps
-) =>
-  JSON.stringify(removeLegendSuppression(currentPlotProps)) !==
-  JSON.stringify(removeLegendSuppression(plotProps))
-
 export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
   plotProps,
   renderZoomedInPlot,
@@ -31,7 +23,9 @@ export const ZoomablePlot: React.FC<ZoomablePlotOwnProps> = ({
 }) => {
   const previousPlotProps = useRef(plotProps)
   useEffect(() => {
-    if (areDifferentPlotProps(previousPlotProps.current, plotProps)) {
+    if (
+      JSON.stringify(previousPlotProps.current) !== JSON.stringify(plotProps)
+    ) {
       renderZoomedInPlot(plotProps, id, true)
       previousPlotProps.current = plotProps
     }
diff --git a/webview/src/plots/components/ZoomedInPlot.tsx b/webview/src/plots/components/ZoomedInPlot.tsx
new file mode 100644
index 0000000000..a0160f18d7
--- /dev/null
+++ b/webview/src/plots/components/ZoomedInPlot.tsx
@@ -0,0 +1,36 @@
+import React from 'react'
+import VegaLite, { VegaLiteProps } from 'react-vega/lib/VegaLite'
+import { Config } from 'vega-lite'
+import merge from 'lodash.merge'
+import cloneDeep from 'lodash.clonedeep'
+import styles from './styles.module.scss'
+import { getThemeValue, ThemeProperty } from '../../util/styles'
+
+export type ZoomedInPlotProps = {
+  props: VegaLiteProps
+}
+
+export const ZoomedInPlot: React.FC<ZoomedInPlotProps> = ({
+  props
+}: ZoomedInPlotProps) => (
+  <div className={styles.zoomedInPlot} data-testid="zoomed-in-plot">
+    <VegaLite
+      {...merge(
+        { ...cloneDeep(props) },
+        {
+          spec: { encoding: { color: { legend: { disable: false } } } }
+        }
+      )}
+      config={{
+        ...(props.config as Config),
+        background: getThemeValue(ThemeProperty.MENU_BACKGROUND)
+      }}
+      actions={{
+        compiled: false,
+        editor: false,
+        export: true,
+        source: false
+      }}
+    />
+  </div>
+)
diff --git a/webview/src/plots/components/util.ts b/webview/src/plots/components/util.ts
index d9055c86fa..96bf5e46bf 100644
--- a/webview/src/plots/components/util.ts
+++ b/webview/src/plots/components/util.ts
@@ -1,5 +1,4 @@
 import { PlotSize } from 'dvc/src/plots/webview/contract'
-import { VegaLiteProps } from 'react-vega/lib/VegaLite'
 
 const MaxItemsBeforeVirtualization = {
   [PlotSize.LARGE]: 10,
@@ -45,19 +44,3 @@ export const getNbItemsPerRow = (size: PlotSize) => {
 
 export const shouldUseVirtualizedGrid = (nbItems: number, size: PlotSize) =>
   nbItems > MaxItemsBeforeVirtualization[size]
-
-export const removeLegendSuppression = (
-  plotProps?: VegaLiteProps
-): VegaLiteProps => {
-  const plot = { ...plotProps } as VegaLiteProps & {
-    spec: {
-      encoding?: {
-        color?: { legend?: { disable?: boolean } }
-      }
-    }
-  }
-  if (plot?.spec?.encoding?.color?.legend?.disable !== undefined) {
-    plot.spec.encoding.color.legend.disable = false
-  }
-  return plot
-}