Fix Comparison Plot "Image By Step" height (#4458)

Author: julieg18

extension/src/test/fixtures/plotsDiff/staticImages/image/7.jpg

@@ -1,5 +1,5 @@
 import { ComparisonPlot } from 'dvc/src/plots/webview/contract'
-import React, { useState } from 'react'
+import React, { useState, useEffect, useRef } from 'react'
 import cx from 'classnames'
 import { useSelector } from 'react-redux'
 import styles from './styles.module.scss'
@@ -27,9 +27,13 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
   nbColumns,
   pinnedColumn
 }) => {
+  const plotsRowRef = useRef<HTMLTableRowElement>(null)
   const draggedId = useSelector(
     (state: PlotsState) => state.dragAndDrop.draggedRef?.itemId
   )
+  const comparisonWidth = useSelector(
+    (state: PlotsState) => state.comparison.width
+  )
   const [isShown, setIsShown] = useState(true)
 
   const toggleIsShownState = () => {
@@ -39,6 +43,37 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
     setIsShown(!isShown)
   }
 
+  const updateMultiImgHeight = (image: HTMLImageElement) => {
+    const aspectRatio = image.naturalWidth / image.naturalHeight
+    const width = image.clientWidth
+    const calculatedHeight = Number.parseFloat((width / aspectRatio).toFixed(2))
+
+    plotsRowRef.current?.style.setProperty(
+      '--img-height',
+      `${calculatedHeight}px`
+    )
+  }
+
+  useEffect(() => {
+    const img: HTMLImageElement | null | undefined =
+      plotsRowRef.current?.querySelector(`.${styles.multiImageWrapper} img`)
+    if (!img) {
+      return
+    }
+
+    if (img.complete) {
+      updateMultiImgHeight(img)
+      return
+    }
+    img.addEventListener(
+      'load',
+      () => {
+        updateMultiImgHeight(img)
+      },
+      { once: true }
+    )
+  }, [comparisonWidth])
+
   return (
     <>
       <tr>
@@ -62,7 +97,7 @@ export const ComparisonTableRow: React.FC<ComparisonTableRowProps> = ({
         </td>
         {nbColumns > 1 && pinnedColumn && <td colSpan={nbColumns - 1}></td>}
       </tr>
-      <tr>
+      <tr ref={plotsRowRef}>
         {plots.map(plot => {
           const isPinned = pinnedColumn === plot.id
           return (

webview/src/plots/components/comparisonTable/ComparisonTableRow.tsx

@@ -176,8 +176,8 @@ $gap: 4px;
 .multiImageWrapper {
   .image,
   .noImageContent {
-    height: 380px;
     object-fit: contain;
+    height: var(--img-height);
   }
 }
 

webview/src/plots/components/comparisonTable/styles.module.scss

@@ -2,7 +2,7 @@ import { configureStore } from '@reduxjs/toolkit'
 import React from 'react'
 import { Provider, useDispatch } from 'react-redux'
 import type { StoryFn, Meta } from '@storybook/react'
-import { userEvent, within } from '@storybook/testing-library'
+import { userEvent, within, fireEvent } from '@storybook/testing-library'
 import {
   PlotsData,
   DEFAULT_SECTION_COLLAPSED,
@@ -227,6 +227,35 @@ AllSmall.args = {
 }
 AllSmall.parameters = CHROMATIC_VIEWPORTS_WITH_DELAY
 
+export const WithMixedMultiImgHeight = Template.bind({})
+WithMixedMultiImgHeight.args = {
+  data: {
+    comparison: {
+      ...comparisonPlotsFixture,
+      plots: comparisonPlotsFixture.plots.filter(({ path }) =>
+        path.includes('image')
+      )
+    }
+  }
+}
+WithMixedMultiImgHeight.play = async ({ canvasElement }) => {
+  const plotSizeSliders = await within(canvasElement).findByTestId(
+    'size-sliders'
+  )
+
+  const sizeSlider = within(plotSizeSliders).getByRole('slider')
+
+  fireEvent.change(sizeSlider, { target: { value: -5 } })
+
+  const multiImgCells = await within(canvasElement).findAllByTestId(
+    'multi-image-cell'
+  )
+
+  const multiImgSlider = within(multiImgCells[1]).getByRole('slider')
+  fireEvent.change(multiImgSlider, { target: { value: 7 } })
+}
+WithMixedMultiImgHeight.parameters = { chromatic: { delay: 2000 } }
+
 export const VirtualizedPlots = Template.bind({})
 VirtualizedPlots.args = {
   data: {