feat(canvas): add custom ref to ScatterPlotCanvas and NetworkCanvas (#1953)

salim-runsafe · web-flow · commit b02107448a17 · 2022-05-10T21:21:16.000+09:00
* Add forwardRef to ScatterPlotCanvas

* Add story

* Adding support for NetworkCanvas

* Aligned function names

* Formatting

* Type fixes
diff --git a/packages/network/src/NetworkCanvas.tsx b/packages/network/src/NetworkCanvas.tsx
@@ -1,4 +1,13 @@
-import { useCallback, useRef, useEffect, createElement, MouseEvent, useMemo } from 'react'
+import {
+    ForwardedRef,
+    forwardRef,
+    useCallback,
+    useRef,
+    useEffect,
+    createElement,
+    MouseEvent,
+    useMemo,
+} from 'react'
 import { getDistance, getRelativeCursor, Container, useDimensions, useTheme } from '@nivo/core'
 import { useTooltip } from '@nivo/tooltip'
 import { useComputedAnnotations, renderAnnotationsToCanvas } from '@nivo/annotations'
@@ -17,7 +26,9 @@ import {
 type InnerNetworkCanvasProps<Node extends InputNode, Link extends InputLink> = Omit<
     NetworkCanvasProps<Node, Link>,
     'renderWrapper' | 'theme'
->
+> & {
+    canvasRef: ForwardedRef<HTMLCanvasElement>
+}
 
 const InnerNetworkCanvas = <Node extends InputNode, Link extends InputLink>({
     width,
@@ -56,6 +67,7 @@ const InnerNetworkCanvas = <Node extends InputNode, Link extends InputLink>({
     defaultActiveNodeIds = canvasDefaultProps.defaultActiveNodeIds,
     nodeTooltip = canvasDefaultProps.nodeTooltip as NodeTooltip<Node>,
     onClick,
+    canvasRef,
 }: InnerNetworkCanvasProps<Node, Link>) => {
     const canvasEl = useRef<HTMLCanvasElement | null>(null)
     const { margin, innerWidth, innerHeight, outerWidth, outerHeight } = useDimensions(
@@ -202,7 +214,10 @@ const InnerNetworkCanvas = <Node extends InputNode, Link extends InputLink>({
 
     return (
         <canvas
-            ref={canvasEl}
+            ref={canvas => {
+                canvasEl.current = canvas
+                if (canvasRef && 'current' in canvasRef) canvasRef.current = canvas
+            }}
             width={outerWidth * pixelRatio}
             height={outerHeight * pixelRatio}
             style={{
@@ -218,18 +233,24 @@ const InnerNetworkCanvas = <Node extends InputNode, Link extends InputLink>({
     )
 }
 
-export const NetworkCanvas = <
-    Node extends InputNode = InputNode,
-    Link extends InputLink = InputLink
->({
-    theme,
-    isInteractive = canvasDefaultProps.isInteractive,
-    animate = canvasDefaultProps.animate,
-    motionConfig = canvasDefaultProps.motionConfig,
-    renderWrapper,
-    ...otherProps
-}: NetworkCanvasProps<Node, Link>) => (
-    <Container {...{ isInteractive, animate, motionConfig, theme, renderWrapper }}>
-        <InnerNetworkCanvas<Node, Link> isInteractive={isInteractive} {...otherProps} />
-    </Container>
+export const NetworkCanvas = forwardRef(
+    <Node extends InputNode = InputNode, Link extends InputLink = InputLink>(
+        {
+            theme,
+            isInteractive = canvasDefaultProps.isInteractive,
+            animate = canvasDefaultProps.animate,
+            motionConfig = canvasDefaultProps.motionConfig,
+            renderWrapper,
+            ...otherProps
+        }: NetworkCanvasProps<Node, Link>,
+        ref: ForwardedRef<HTMLCanvasElement>
+    ) => (
+        <Container {...{ isInteractive, animate, motionConfig, theme, renderWrapper }}>
+            <InnerNetworkCanvas<Node, Link>
+                isInteractive={isInteractive}
+                {...otherProps}
+                canvasRef={ref}
+            />
+        </Container>
+    )
 )
diff --git a/packages/network/src/ResponsiveNetworkCanvas.tsx b/packages/network/src/ResponsiveNetworkCanvas.tsx
@@ -1,16 +1,28 @@
 import { ResponsiveWrapper } from '@nivo/core'
+import { ForwardedRef, forwardRef } from 'react'
 import { NetworkCanvasProps, InputNode, InputLink } from './types'
 import { NetworkCanvas } from './NetworkCanvas'
 
-export const ResponsiveNetworkCanvas = <
+export const ResponsiveNetworkCanvas = forwardRef(function ResponsiveBarCanvas<
     Node extends InputNode = InputNode,
     Link extends InputLink = InputLink
 >(
-    props: Omit<NetworkCanvasProps<Node, Link>, 'height' | 'width'>
-) => (
-    <ResponsiveWrapper>
-        {({ width, height }) => (
-            <NetworkCanvas<Node, Link> width={width} height={height} {...props} />
-        )}
-    </ResponsiveWrapper>
-)
+    props: Omit<NetworkCanvasProps<Node, Link>, 'height' | 'width'>,
+    ref: ForwardedRef<HTMLCanvasElement>
+) {
+    return (
+        <ResponsiveWrapper>
+            {({ width, height }) => (
+                <NetworkCanvas
+                    width={width}
+                    height={height}
+                    {...(props as Omit<
+                        NetworkCanvasProps<InputNode, InputLink>,
+                        'height' | 'width'
+                    >)}
+                    ref={ref}
+                />
+            )}
+        </ResponsiveWrapper>
+    )
+})
diff --git a/packages/network/stories/networkCanvas.stories.tsx b/packages/network/stories/networkCanvas.stories.tsx
@@ -9,6 +9,7 @@ import {
     NodeTooltipProps,
     // @ts-ignore
 } from '../src'
+import { useRef } from 'react'
 
 export default {
     component: NetworkCanvas,
@@ -70,3 +71,22 @@ export const CustomNodeRenderer = () => (
 export const OnClickHandler = () => (
     <NetworkCanvas<Node, Link> {...commonProperties} onClick={action('onClick')} />
 )
+
+export const CustomCanvasRef = () => {
+    const ref = useRef(undefined)
+
+    const download = ref => {
+        const canvas = ref.current
+        const link = document.createElement('a')
+        link.download = 'test.png'
+        link.href = canvas.toDataURL('image/png')
+        link.click()
+    }
+
+    return (
+        <>
+            <NetworkCanvas<Node, Link> {...commonProperties} ref={ref} />
+            <button onClick={() => download(ref)}>Download PNG</button>
+        </>
+    )
+}
diff --git a/packages/scatterplot/src/ResponsiveScatterPlotCanvas.tsx b/packages/scatterplot/src/ResponsiveScatterPlotCanvas.tsx
@@ -1,13 +1,28 @@
 import { ResponsiveWrapper } from '@nivo/core'
+import { ForwardedRef, forwardRef } from 'react'
+
 import { ScatterPlotCanvas } from './ScatterPlotCanvas'
 import { ScatterPlotCanvasProps, ScatterPlotDatum } from './types'
 
-export const ResponsiveScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>(
-    props: Omit<ScatterPlotCanvasProps<RawDatum>, 'width' | 'height'>
-) => (
-    <ResponsiveWrapper>
-        {({ width, height }) => (
-            <ScatterPlotCanvas<RawDatum> width={width} height={height} {...props} />
-        )}
-    </ResponsiveWrapper>
-)
+export const ResponsiveScatterPlotCanvas = forwardRef(function ResponsiveScatterPlotCanvas<
+    RawDatum extends ScatterPlotDatum
+>(
+    props: Omit<ScatterPlotCanvasProps<RawDatum>, 'width' | 'height'>,
+    ref: ForwardedRef<HTMLCanvasElement>
+) {
+    return (
+        <ResponsiveWrapper>
+            {({ width, height }) => (
+                <ScatterPlotCanvas
+                    width={width}
+                    height={height}
+                    {...(props as Omit<
+                        ScatterPlotCanvasProps<ScatterPlotDatum>,
+                        'height' | 'width'
+                    >)}
+                    ref={ref}
+                />
+            )}
+        </ResponsiveWrapper>
+    )
+})
diff --git a/packages/scatterplot/src/ScatterPlotCanvas.tsx b/packages/scatterplot/src/ScatterPlotCanvas.tsx
@@ -1,4 +1,13 @@
-import { createElement, useRef, useState, useEffect, useCallback, useMemo } from 'react'
+import {
+    ForwardedRef,
+    createElement,
+    forwardRef,
+    useCallback,
+    useEffect,
+    useMemo,
+    useRef,
+    useState,
+} from 'react'
 import { Container, useDimensions, useTheme, getRelativeCursor, isCursorInRect } from '@nivo/core'
 import { renderAnnotationsToCanvas } from '@nivo/annotations'
 import { CanvasAxisProps, renderAxesToCanvas, renderGridLinesToCanvas } from '@nivo/axes'
@@ -12,7 +21,9 @@ import { ScatterPlotCanvasProps, ScatterPlotDatum, ScatterPlotNodeData } from '.
 type InnerScatterPlotCanvasProps<RawDatum extends ScatterPlotDatum> = Omit<
     ScatterPlotCanvasProps<RawDatum>,
     'renderWrapper' | 'theme'
->
+> & {
+    canvasRef: ForwardedRef<HTMLCanvasElement>
+}
 
 const InnerScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>({
     data,
@@ -46,6 +57,7 @@ const InnerScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>({
     onClick,
     tooltip = canvasDefaultProps.tooltip,
     legends = canvasDefaultProps.legends,
+    canvasRef,
 }: InnerScatterPlotCanvasProps<RawDatum>) => {
     const canvasEl = useRef<HTMLCanvasElement | null>(null)
     const theme = useTheme()
@@ -270,7 +282,10 @@ const InnerScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>({
 
     return (
         <canvas
-            ref={canvasEl}
+            ref={canvas => {
+                canvasEl.current = canvas
+                if (canvasRef && 'current' in canvasRef) canvasRef.current = canvas
+            }}
             width={outerWidth * pixelRatio}
             height={outerHeight * pixelRatio}
             style={{
@@ -286,13 +301,13 @@ const InnerScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>({
     )
 }
 
-export const ScatterPlotCanvas = <RawDatum extends ScatterPlotDatum>({
-    isInteractive,
-    renderWrapper,
-    theme,
-    ...props
-}: ScatterPlotCanvasProps<RawDatum>) => (
-    <Container {...{ isInteractive, renderWrapper, theme }} animate={false}>
-        <InnerScatterPlotCanvas<RawDatum> {...props} />
-    </Container>
+export const ScatterPlotCanvas = forwardRef(
+    <RawDatum extends ScatterPlotDatum>(
+        { isInteractive, renderWrapper, theme, ...props }: ScatterPlotCanvasProps<RawDatum>,
+        ref: ForwardedRef<HTMLCanvasElement>
+    ) => (
+        <Container {...{ isInteractive, renderWrapper, theme }} animate={false}>
+            <InnerScatterPlotCanvas<RawDatum> {...props} canvasRef={ref} />
+        </Container>
+    )
 )
diff --git a/packages/scatterplot/stories/ScatterPlotCanvas.stories.tsx b/packages/scatterplot/stories/ScatterPlotCanvas.stories.tsx
@@ -1,8 +1,8 @@
-import { useState, useCallback, useMemo } from 'react'
+import { useCallback, useMemo, useRef, useState } from 'react'
 import omit from 'lodash/omit'
 import { Meta } from '@storybook/react'
 // @ts-ignore
-import { ScatterPlotCanvas, ResponsiveScatterPlotCanvas, ScatterPlotNodeData } from '../src'
+import { ResponsiveScatterPlotCanvas, ScatterPlotCanvas, ScatterPlotNodeData } from '../src'
 
 export default {
     component: ScatterPlotCanvas,
@@ -403,3 +403,42 @@ export const CustomTooltip = () => (
         )}
     />
 )
+
+export const CustomCanvasRef = () => {
+    const ref = useRef(undefined)
+
+    const download = ref => {
+        const canvas = ref.current
+        const link = document.createElement('a')
+        link.download = 'test.png'
+        link.href = canvas.toDataURL('image/png')
+        link.click()
+    }
+
+    return (
+        <>
+            <ScatterPlotCanvas<SampleDatum>
+                {...commonProps}
+                ref={ref}
+                tooltip={({ node }) => (
+                    <div
+                        style={{
+                            color: node.color,
+                            background: '#333',
+                            padding: '12px 16px',
+                        }}
+                    >
+                        <strong>
+                            {node.id} ({node.serieId})
+                        </strong>
+                        <br />
+                        {`x: ${node.formattedX}`}
+                        <br />
+                        {`y: ${node.formattedY}`}
+                    </div>
+                )}
+            />
+            <button onClick={() => download(ref)}>Download PNG</button>
+        </>
+    )
+}
