From c87658e58a563822f2364ce4f93dc5dc32bb88b9 Mon Sep 17 00:00:00 2001
From: plouc <benitteraphael@gmail.com>
Date: Sun, 19 Nov 2023 11:38:08 +0900
Subject: [PATCH] feat(pie): add the ability to programmatically control the
 activeId

---
 packages/pie/src/Pie.tsx              |  4 ++
 packages/pie/src/hooks.ts             | 73 +++++++++++++++++++++++++--
 packages/pie/src/types.ts             |  3 ++
 storybook/stories/pie/Pie.stories.tsx | 36 +++++++++++++
 4 files changed, 111 insertions(+), 5 deletions(-)

diff --git a/packages/pie/src/Pie.tsx b/packages/pie/src/Pie.tsx
index a6189c455..b37a0b842 100644
--- a/packages/pie/src/Pie.tsx
+++ b/packages/pie/src/Pie.tsx
@@ -74,6 +74,8 @@ const InnerPie = <RawDatum extends MayHaveLabel>({
     onMouseMove,
     onMouseLeave,
     tooltip = defaultProps.tooltip,
+    activeId: activeIdFromProps,
+    onActiveIdChange,
 
     transitionMode = defaultProps.transitionMode,
 
@@ -117,6 +119,8 @@ const InnerPie = <RawDatum extends MayHaveLabel>({
         cornerRadius,
         activeInnerRadiusOffset,
         activeOuterRadiusOffset,
+        activeId: activeIdFromProps,
+        onActiveIdChange,
     })
 
     const boundDefs = bindDefs(defs, dataWithArc, fill)
diff --git a/packages/pie/src/hooks.ts b/packages/pie/src/hooks.ts
index 83be94273..a640d937a 100644
--- a/packages/pie/src/hooks.ts
+++ b/packages/pie/src/hooks.ts
@@ -165,6 +165,45 @@ export const usePieArcs = <RawDatum>({
     ])
 }
 
+/**
+ * Encapsulate the logic for defining/reading the active arc ID,
+ * which can be either controlled (handled externally), or uncontrolled
+ * (handled internally), we can optionally define a default value when
+ * it's uncontrolled.
+ */
+const useActiveId = ({
+    activeId: activeIdFromProps,
+    onActiveIdChange,
+    defaultActiveId = null,
+}: {
+    activeId?: DatumId | null
+    onActiveIdChange?: (id: DatumId | null) => void
+    defaultActiveId?: DatumId | null
+}) => {
+    const isControlled = typeof activeIdFromProps != 'undefined'
+
+    const [internalActiveId, setInternalActiveId] = useState<DatumId | null>(
+        !isControlled ? defaultActiveId : null
+    )
+
+    const activeId = isControlled ? activeIdFromProps : internalActiveId
+
+    const setActiveId = useCallback(
+        (id: DatumId | null) => {
+            if (onActiveIdChange) {
+                onActiveIdChange(id)
+            }
+
+            if (!isControlled) {
+                setInternalActiveId(id)
+            }
+        },
+        [isControlled, onActiveIdChange, setInternalActiveId]
+    )
+
+    return { activeId, setActiveId }
+}
+
 /**
  * Compute pie layout using explicit radius/innerRadius,
  * expressed in pixels.
@@ -180,6 +219,9 @@ export const usePie = <RawDatum>({
     cornerRadius = defaultProps.cornerRadius,
     activeInnerRadiusOffset = defaultProps.activeInnerRadiusOffset,
     activeOuterRadiusOffset = defaultProps.activeOuterRadiusOffset,
+    activeId: activeIdFromProps,
+    onActiveIdChange,
+    defaultActiveId,
 }: Pick<
     Partial<CompletePieSvgProps<RawDatum>>,
     | 'startAngle'
@@ -189,12 +231,20 @@ export const usePie = <RawDatum>({
     | 'cornerRadius'
     | 'activeInnerRadiusOffset'
     | 'activeOuterRadiusOffset'
+    | 'activeId'
+    | 'onActiveIdChange'
+    | 'defaultActiveId'
 > & {
     data: Omit<ComputedDatum<RawDatum>, 'arc'>[]
     radius: number
     innerRadius: number
 }) => {
-    const [activeId, setActiveId] = useState<DatumId | null>(null)
+    const { activeId, setActiveId } = useActiveId({
+        activeId: activeIdFromProps,
+        onActiveIdChange,
+        defaultActiveId,
+    })
+
     const [hiddenIds, setHiddenIds] = useState<DatumId[]>([])
     const pieArcs = usePieArcs({
         data,
@@ -242,6 +292,9 @@ export const usePieFromBox = <RawDatum>({
     fit = defaultProps.fit,
     activeInnerRadiusOffset = defaultProps.activeInnerRadiusOffset,
     activeOuterRadiusOffset = defaultProps.activeOuterRadiusOffset,
+    activeId: activeIdFromProps,
+    onActiveIdChange,
+    defaultActiveId,
 }: Pick<
     CompletePieSvgProps<RawDatum>,
     | 'width'
@@ -255,10 +308,19 @@ export const usePieFromBox = <RawDatum>({
     | 'fit'
     | 'activeInnerRadiusOffset'
     | 'activeOuterRadiusOffset'
-> & {
-    data: Omit<ComputedDatum<RawDatum>, 'arc'>[]
-}) => {
-    const [activeId, setActiveId] = useState<string | number | null>(null)
+> &
+    Pick<
+        Partial<CompletePieSvgProps<RawDatum>>,
+        'activeId' | 'onActiveIdChange' | 'defaultActiveId'
+    > & {
+        data: Omit<ComputedDatum<RawDatum>, 'arc'>[]
+    }) => {
+    const { activeId, setActiveId } = useActiveId({
+        activeId: activeIdFromProps,
+        onActiveIdChange,
+        defaultActiveId,
+    })
+
     const [hiddenIds, setHiddenIds] = useState<DatumId[]>([])
     const computedProps = useMemo(() => {
         let radius = Math.min(width, height) / 2
@@ -335,6 +397,7 @@ export const usePieFromBox = <RawDatum>({
 
     return {
         arcGenerator,
+        activeId,
         setActiveId,
         toggleSerie,
         ...pieArcs,
diff --git a/packages/pie/src/types.ts b/packages/pie/src/types.ts
index 66f926e8b..22401fbf1 100644
--- a/packages/pie/src/types.ts
+++ b/packages/pie/src/types.ts
@@ -113,6 +113,9 @@ export type CommonPieProps<RawDatum> = {
     // interactivity
     isInteractive: boolean
     tooltip: React.FC<PieTooltipProps<RawDatum>>
+    activeId: DatumId | null
+    onActiveIdChange: (id: DatumId | null) => void
+    defaultActiveId: DatumId | null
 
     legends: readonly LegendProps[]
 
diff --git a/storybook/stories/pie/Pie.stories.tsx b/storybook/stories/pie/Pie.stories.tsx
index 38750b3d7..10748a03e 100644
--- a/storybook/stories/pie/Pie.stories.tsx
+++ b/storybook/stories/pie/Pie.stories.tsx
@@ -1,3 +1,4 @@
+import { useState } from 'react'
 import type { Meta, StoryObj } from '@storybook/react'
 import { animated } from '@react-spring/web'
 import { generateProgrammingLanguageStats } from '@nivo/generators'
@@ -245,3 +246,38 @@ export const CustomArcLabelComponent: Story = {
         />
     ),
 }
+
+const controlledPieProps = {
+    ...commonProperties,
+    width: 400,
+    height: 400,
+    margin: { top: 60, right: 80, bottom: 60, left: 80 },
+    innerRadius: 0.4,
+    padAngle: 0.3,
+    cornerRadius: 3,
+    activeOuterRadiusOffset: 12,
+    activeInnerRadiusOffset: 12,
+    arcLinkLabelsDiagonalLength: 10,
+    arcLinkLabelsStraightLength: 10,
+}
+
+const ControlledPies = () => {
+    const [activeId, setActiveId] = useState<string>(commonProperties.data[1].id)
+
+    return (
+        <div
+            style={{
+                width: '800px',
+                display: 'grid',
+                gridTemplateColumns: 'repeat(2, 1fr)',
+            }}
+        >
+            <Pie {...controlledPieProps} activeId={activeId} onActiveIdChange={setActiveId} />
+            <Pie {...controlledPieProps} activeId={activeId} onActiveIdChange={setActiveId} />
+        </div>
+    )
+}
+
+export const ControlledActiveId: Story = {
+    render: () => <ControlledPies />,
+}