Use composition for annotations with chart context

Author: adamwoodnz

projects/js-packages/charts/src/components/line-chart/line-chart-annotations-overlay.tsx

@@ -1,30 +1,32 @@
 import { DataContext } from '@visx/xychart';
 import { useEffect, useState, useCallback } from 'react';
-import LineChartAnnotation from './line-chart-annotation';
+import { useChartContext } from '../../providers/chart-context';
 import styles from './line-chart.module.scss';
-import type { LineChartRef } from './line-chart';
-import type { LineChartAnnotationProps } from './line-chart-annotation';
 import type { AxisScale } from '@visx/axis';
-import type { FC, RefObject } from 'react';
+import type { FC, ReactNode } from 'react';
 
 interface LineChartAnnotationsProps {
-	chartRef: RefObject< LineChartRef >;
-	annotations: LineChartAnnotationProps[];
-	chartWidth: number;
-	chartHeight: number;
+	chartId: string;
+	children?: ReactNode;
 }
 
 interface ScaleData {
 	xScale: AxisScale< Date >;
 	yScale: AxisScale< number >;
 }
 
-const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( {
-	chartRef,
-	annotations,
-	chartWidth,
-	chartHeight,
-} ) => {
+const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( { chartId, children } ) => {
+	const { getChartData } = useChartContext();
+
+	let chartData = null;
+	if ( chartId ) {
+		chartData = getChartData( chartId );
+	}
+
+	const chartRef = chartData?.chartRef;
+	const chartWidth = chartData?.chartWidth || 0;
+	const chartHeight = chartData?.chartHeight || 0;
+
 	const [ scales, setScales ] = useState< ScaleData | null >( null );
 	const [ scalesStable, setScalesStable ] = useState< boolean >( false );
 
@@ -42,7 +44,7 @@ const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( {
 
 	// Get scales from chart ref and return them with signature for comparison
 	const getScalesData = useCallback( () => {
-		if ( chartRef.current ) {
+		if ( chartRef?.current ) {
 			const scaleData = chartRef.current.getScales();
 
 			if ( scaleData ) {
@@ -107,6 +109,11 @@ const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( {
 		};
 	}, [ getScalesData, chartWidth, chartHeight ] );
 
+	// Early return if no chart data available
+	if ( ! chartRef || ! children ) {
+		return null;
+	}
+
 	if ( ! scales || ! scalesStable ) {
 		return null;
 	}
@@ -116,16 +123,9 @@ const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( {
 	const dataContextValue = {
 		xScale: scales.xScale,
 		yScale: scales.yScale,
-		// Add minimal required properties for DataContext
-		theme: {},
 		margin: { top: 0, right: 0, bottom: 0, left: 0 },
 		width: chartWidth,
 		height: chartHeight,
-		// Additional visx DataContext properties that may be needed
-		colorScale: undefined,
-		dataRegistry: new Map(),
-		registerData: () => {},
-		unregisterData: () => {},
 	} as unknown as Parameters< typeof DataContext.Provider >[ 0 ][ 'value' ];
 
 	return (
@@ -136,13 +136,7 @@ const LineChartAnnotationsOverlay: FC< LineChartAnnotationsProps > = ( {
 				className={ styles[ 'line-chart__annotations-overlay' ] }
 				data-testid="line-chart-annotations-overlay"
 			>
-				{ annotations.map( ( annotation, index ) => (
-					<LineChartAnnotation
-						key={ `overlay-annotation-${ index }` }
-						testId={ `overlay-annotation-${ index }` }
-						{ ...annotation }
-					/>
-				) ) }
+				{ children }
 			</svg>
 		</DataContext.Provider>
 	);

projects/js-packages/charts/src/components/line-chart/line-chart.tsx

@@ -21,6 +21,9 @@ import {
 	useState,
 	useRef,
 	useCallback,
+	Children,
+	cloneElement,
+	isValidElement,
 } from 'react';
 import { ChartProvider, useChartId, useChartRegistration } from '../../providers/chart-context';
 import { useXYChartTheme, useChartTheme } from '../../providers/theme/theme-provider';
@@ -30,9 +33,7 @@ import { useChartDataTransform } from '../shared/use-chart-data-transform';
 import { useChartMargin } from '../shared/use-chart-margin';
 import { useElementHeight } from '../shared/use-element-height';
 import { withResponsive } from '../shared/with-responsive';
-import LineChartAnnotationsOverlay from './line-chart-annotations-overlay';
 import styles from './line-chart.module.scss';
-import type { LineChartAnnotationProps } from './line-chart-annotation';
 import type { BaseChartProps, DataPoint, DataPointDate, SeriesData } from '../../types';
 import type { TickFormatter } from '@visx/axis';
 import type { GlyphProps } from '@visx/xychart';
@@ -143,7 +144,7 @@ interface LineChartProps extends BaseChartProps< SeriesData[] > {
 		showVertical?: boolean;
 		showHorizontal?: boolean;
 	};
-	annotations?: LineChartAnnotationProps[];
+	children?: ReactNode;
 }
 
 type TooltipDatum = {
@@ -294,11 +295,11 @@ const LineChartInternal = forwardRef< LineChartRef, LineChartProps >(
 			renderTooltip = renderDefaultTooltip,
 			withStartGlyphs = false,
 			options = {},
-			annotations,
 			onPointerDown = undefined,
 			onPointerUp = undefined,
 			onPointerMove = undefined,
 			onPointerOut = undefined,
+			children,
 		},
 		ref
 	) => {
@@ -404,13 +405,25 @@ const LineChartInternal = forwardRef< LineChartRef, LineChartProps >(
 		);
 
 		// Register chart with context only if data is valid
-		useChartRegistration( chartId, legendItems, providerTheme, 'line', isDataValid, {
-			withGradientFill,
-			smoothing,
-			curveType,
-			withStartGlyphs,
-			withLegendGlyph,
-		} );
+		useChartRegistration(
+			chartId,
+			legendItems,
+			providerTheme,
+			'line',
+			isDataValid,
+			{
+				withGradientFill,
+				smoothing,
+				curveType,
+				withStartGlyphs,
+				withLegendGlyph,
+			},
+			{
+				chartRef: internalChartRef,
+				chartWidth: width,
+				chartHeight: height - ( showLegend ? legendHeight : 0 ),
+			}
+		);
 
 		const accessors = {
 			xAccessor: ( d: DataPointDate ) => d?.date,
@@ -632,16 +645,6 @@ const LineChartInternal = forwardRef< LineChartRef, LineChartProps >(
 					</XYChart>
 				</div>
 
-				{ /* Render annotations as external overlay to avoid interaction blocking */ }
-				{ annotations?.length && (
-					<LineChartAnnotationsOverlay
-						chartRef={ internalChartRef }
-						annotations={ annotations }
-						chartWidth={ width }
-						chartHeight={ height - ( showLegend ? legendHeight : 0 ) }
-					/>
-				) }
-
 				{ showLegend && (
 					<Legend
 						items={ legendItems }
@@ -653,6 +656,15 @@ const LineChartInternal = forwardRef< LineChartRef, LineChartProps >(
 						ref={ legendRef }
 					/>
 				) }
+
+				{ /* Render children here so they have access to the chart ID */ }
+				{ Children.map( children, child => {
+					if ( isValidElement( child ) ) {
+						return cloneElement( child, { chartId } as { chartId: string } );
+					}
+
+					return child;
+				} ) }
 			</div>
 		);
 	}