feat: Allow ref callback for Chart and ChartPanel (#2174)

Needed for deephaven/deephaven-plugins#657
because when using ChartPanel `containerRef` is undefined on first
render due to the loading logic. This means we need to allow a ref
callback so we can use a callback to set state in the plugin and add
event listeners to the container when it actually exists

packages/chart/src/Chart.tsx

@@ -28,6 +28,7 @@ import {
   ModeBarButtonAny,
 } from 'plotly.js';
 import type { PlotParams } from 'react-plotly.js';
+import { mergeRefs } from '@deephaven/react-hooks';
 import { bindAllMethods } from '@deephaven/utils';
 import createPlotlyComponent from './plotly/createPlotlyComponent';
 import Plotly from './plotly/Plotly';
@@ -57,7 +58,7 @@ interface ChartProps {
 
   isActive: boolean;
   Plotly: typeof Plotly;
-  containerRef?: React.RefObject<HTMLDivElement>;
+  containerRef?: React.Ref<HTMLDivElement>;
   onDisconnect: () => void;
   onReconnect: () => void;
   onUpdate: (obj: { isLoading: boolean }) => void;
@@ -156,7 +157,8 @@ class Chart extends Component<ChartProps, ChartState> {
 
     this.PlotComponent = createPlotlyComponent(props.Plotly);
     this.plot = React.createRef();
-    this.plotWrapper = props.containerRef ?? React.createRef();
+    this.plotWrapper = React.createRef();
+    this.plotWrapperMerged = mergeRefs(this.plotWrapper, props.containerRef);
     this.columnFormats = [];
     this.dateTimeFormatterOptions = {};
     this.decimalFormatOptions = {};
@@ -238,6 +240,8 @@ class Chart extends Component<ChartProps, ChartState> {
 
   plotWrapper: RefObject<HTMLDivElement>;
 
+  plotWrapperMerged: React.RefCallback<HTMLDivElement>;
+
   columnFormats?: FormattingRule[];
 
   dateTimeFormatterOptions?: DateTimeColumnFormatterOptions;
@@ -713,7 +717,7 @@ class Chart extends Component<ChartProps, ChartState> {
     const isPlotShown = data != null;
 
     return (
-      <div className="h-100 w-100 chart-wrapper" ref={this.plotWrapper}>
+      <div className="h-100 w-100 chart-wrapper" ref={this.plotWrapperMerged}>
         {isPlotShown && (
           <PlotComponent
             // eslint-disable-next-line @typescript-eslint/ban-ts-comment

packages/components/src/spectrum/comboBox/ComboBox.tsx

@@ -5,9 +5,9 @@ import {
 } from '@adobe/react-spectrum';
 import type { DOMRef } from '@react-types/shared';
 import cl from 'classnames';
+import { useMergeRef } from '@deephaven/react-hooks';
 import type { NormalizedItem } from '../utils';
 import { PickerPropsT, usePickerProps } from '../picker';
-import useMultiRef from '../picker/useMultiRef';
 
 export type ComboBoxProps = PickerPropsT<SpectrumComboBoxProps<NormalizedItem>>;
 
@@ -22,7 +22,7 @@ export const ComboBox = React.forwardRef(function ComboBox(
     ref: scrollRef,
     ...comboBoxProps
   } = usePickerProps(props);
-  const pickerRef = useMultiRef(ref, scrollRef);
+  const pickerRef = useMergeRef(ref, scrollRef);
   return (
     <SpectrumComboBox
       // eslint-disable-next-line react/jsx-props-no-spreading

packages/components/src/spectrum/picker/Picker.tsx

@@ -5,9 +5,9 @@ import {
 import type { DOMRef } from '@react-types/shared';
 import cl from 'classnames';
 import React from 'react';
+import { useMergeRef } from '@deephaven/react-hooks';
 import type { NormalizedItem } from '../utils';
 import type { PickerProps } from './PickerProps';
-import useMultiRef from './useMultiRef';
 import { usePickerProps } from './usePickerProps';
 
 /**
@@ -28,7 +28,7 @@ export const Picker = React.forwardRef(function Picker(
     ref: scrollRef,
     ...pickerProps
   } = usePickerProps(props);
-  const pickerRef = useMultiRef(ref, scrollRef);
+  const pickerRef = useMergeRef(ref, scrollRef);
   return (
     <SpectrumPicker
       // eslint-disable-next-line react/jsx-props-no-spreading

packages/dashboard-core-plugins/src/panels/ChartPanel.tsx

@@ -123,8 +123,11 @@ interface OwnProps extends DashboardPanelProps {
   makeModel: () => Promise<ChartModel>;
   localDashboardId: string;
   Plotly?: typeof PlotlyType;
-  /** The plot container div */
-  containerRef?: RefObject<HTMLDivElement>;
+  /**
+   * The plot container div.
+   * The ref will be undefined on initial render if the chart needs to be loaded.
+   */
+  containerRef?: React.Ref<HTMLDivElement>;
 
   panelState?: GLChartPanelState;
 }

packages/react-hooks/src/index.ts

@@ -34,3 +34,4 @@ export * from './useSetAttributesCallback';
 export * from './useSpectrumDisableSpellcheckRef';
 export * from './useWindowedListData';
 export * from './useResizeObserver';
+export * from './useMergeRef';

packages/components/src/spectrum/picker/useMultiRef.test.ts → packages/react-hooks/src/useMergeRef.test.ts

@@ -1,12 +1,52 @@
+import React from 'react';
 import { renderHook } from '@testing-library/react-hooks';
-import useMultiRef from './useMultiRef';
+import { useMergeRef, mergeRefs } from './useMergeRef';
 
-describe('useMultiRef', () => {
+describe('mergeRefs', () => {
+  it('merges ref objects', () => {
+    const refA = React.createRef();
+    const refB = React.createRef();
+    const mergedRef = mergeRefs(refA, refB);
+
+    const refValue = {};
+    mergedRef(refValue);
+    expect(refA.current).toBe(refValue);
+    expect(refB.current).toBe(refValue);
+  });
+
+  it('merges ref callbacks', () => {
+    const refA: React.RefCallback<unknown> = jest.fn();
+    const refB: React.RefCallback<unknown> = jest.fn();
+    const mergedRef = mergeRefs(refA, refB);
+
+    const refValue = {};
+    mergedRef(refValue);
+    expect(refA).toHaveBeenCalledWith(refValue);
+    expect(refB).toHaveBeenCalledWith(refValue);
+  });
+
+  it('ignores null/undefined refs', () => {
+    const refA = React.createRef();
+    const refB: React.RefCallback<unknown> = jest.fn();
+    const refC = null;
+    const refD = undefined;
+    const mergedRef = mergeRefs(refA, refB, refC, refD);
+
+    const refValue = {};
+    mergedRef(refValue);
+    expect(refA.current).toBe(refValue);
+    expect(refB).toHaveBeenCalledWith(refValue);
+    expect(refC).toBe(null);
+    expect(refD).toBe(undefined);
+  });
+});
+
+describe('useMergeRef', () => {
   it('should assign the ref to all refs passed in', () => {
     const ref1 = jest.fn();
     const ref2 = jest.fn();
     const ref3 = jest.fn();
-    const { result } = renderHook(() => useMultiRef(ref1, ref2, ref3));
+    const { result } = renderHook(() => useMergeRef(ref1, ref2, ref3));
     const multiRef = result.current;
     const element = document.createElement('div');
     multiRef(element);
@@ -19,7 +59,7 @@ describe('useMultiRef', () => {
     const ref1 = jest.fn();
     const ref2 = jest.fn();
     const ref3 = jest.fn();
-    const { result } = renderHook(() => useMultiRef(ref1, ref2, ref3));
+    const { result } = renderHook(() => useMergeRef(ref1, ref2, ref3));
     const multiRef = result.current;
     multiRef(null);
     expect(ref1).toHaveBeenCalledWith(null);
@@ -32,7 +72,7 @@ describe('useMultiRef', () => {
     const ref2 = { current: null };
     const ref3 = { current: null };
     const { result } = renderHook(() =>
-      useMultiRef<HTMLDivElement | null>(ref1, ref2, ref3)
+      useMergeRef<HTMLDivElement | null>(ref1, ref2, ref3)
     );
     const multiRef = result.current;
     const element = document.createElement('div');
@@ -47,7 +87,7 @@ describe('useMultiRef', () => {
     const ref2 = { current: null };
     const ref3 = jest.fn();
     const { result } = renderHook(() =>
-      useMultiRef<HTMLDivElement | null>(ref1, ref2, ref3)
+      useMergeRef<HTMLDivElement | null>(ref1, ref2, ref3)
     );
     const multiRef = result.current;
     const element = document.createElement('div');

packages/react-hooks/src/useMergeRef.ts

@@ -0,0 +1,43 @@
+import {
+  type LegacyRef,
+  type MutableRefObject,
+  type Ref,
+  type RefCallback,
+  useMemo,
+} from 'react';
+
+/**
+ * Merge multiple react refs into a single ref callback.
+ * This can be used to merge callback and object refs into a single ref.
+ * Merged callback refs will be called while object refs will have their current property set.
+ * @param refs The refs to merge
+ * @returns A ref callback that will set the value on all refs
+ */
+export function mergeRefs<T = unknown>(
+  ...refs: readonly (MutableRefObject<T> | LegacyRef<T> | null | undefined)[]
+): RefCallback<T> {
+  return newRef => {
+    refs.forEach(ref => {
+      if (ref != null) {
+        if (typeof ref === 'function') {
+          ref(newRef);
+        } else {
+          // React marks RefObject as readonly, but it's just to indicate React manages it
+          // We can still write to its current value
+          // eslint-disable-next-line no-param-reassign
+          (ref as React.MutableRefObject<T | null>).current = newRef;
+        }
+      }
+    });
+  };
+}
+
+/**
+ * Merges multiple refs into one ref that can be assigned to the component.
+ * In turn all the refs passed in will be assigned when the ref returned is assigned.
+ * @param refs Array of refs to assign
+ */
+export function useMergeRef<T>(...refs: readonly Ref<T>[]): RefCallback<T> {
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  return useMemo(() => mergeRefs(...refs), refs);
+}