From 855313c3076fae304e70a2e51898d9f289ff0ccf Mon Sep 17 00:00:00 2001
From: Chris Williams <chris.williams@airbnb.com>
Date: Wed, 12 Aug 2020 11:35:58 -0700
Subject: [PATCH] new(xy-chart): add hook to context for full scale
 modification, not just domain

---
 .../src/sandboxes/vx-chart-poc/Example.tsx    | 37 +++++++++--------
 .../components/providers/ChartProvider.tsx    | 41 +++++++++----------
 .../src/components/series/Stack.tsx           | 32 ++++++++++++---
 .../vx-chart-poc/src/types/context.ts         | 14 ++++---
 4 files changed, 75 insertions(+), 49 deletions(-)

diff --git a/packages/vx-demo/src/sandboxes/vx-chart-poc/Example.tsx b/packages/vx-demo/src/sandboxes/vx-chart-poc/Example.tsx
index 7b282cf13..bf8f2a260 100644
--- a/packages/vx-demo/src/sandboxes/vx-chart-poc/Example.tsx
+++ b/packages/vx-demo/src/sandboxes/vx-chart-poc/Example.tsx
@@ -20,13 +20,7 @@ import Stack from './src/components/series/Stack';
 
 type DataKeys = 'austin' | 'sf' | 'ny';
 
-const data = cityTemperature.slice(100, 100 + 12).map(({ date, ...d }) => ({
-  ...d,
-  // current format is like `20200105` which you can't form a valid date from
-  // @TODO PR soon!
-  date: `${date.slice(0, 4)}-${date.slice(4, 6)}-${date.slice(6, 8)}`,
-  key: date,
-})) as CityTemperature[];
+const data = cityTemperature.slice(100, 100 + 16);
 
 // @TODO wip updating data, not currently used
 // const halfData = data.slice(0, Math.floor(data.length / 2));
@@ -129,7 +123,17 @@ export default function Example() {
   const [renderTooltipInPortal, setRenderTooltipInPortal] = useState(true);
   const [visibleSeries, setVisibleSeries] = useState<
     ('line' | 'bar' | 'groupedbar' | 'stackedbar')[]
-  >(['stackedbar']);
+  >(['bar']);
+  const canSnapTooltipToDataX =
+    (visibleSeries.includes('groupedbar') && renderHorizontally) ||
+    (visibleSeries.includes('stackedbar') && !renderHorizontally) ||
+    visibleSeries.includes('bar');
+
+  const canSnapTooltipToDataY =
+    (visibleSeries.includes('groupedbar') && !renderHorizontally) ||
+    (visibleSeries.includes('stackedbar') && renderHorizontally) ||
+    visibleSeries.includes('bar');
+
   const dateScaleConfig: ScaleConfig<string> = useMemo(() => ({ type: 'band', padding: 0.2 }), []);
   const temperatureScaleConfig: ScaleConfig<number> = useMemo(
     () => ({
@@ -262,8 +266,8 @@ export default function Example() {
               />
             </XYChart>
             <Tooltip
-              snapToDataX={snapTooltipToDataX}
-              snapToDataY={snapTooltipToDataY}
+              snapToDataX={snapTooltipToDataX && canSnapTooltipToDataX}
+              snapToDataY={snapTooltipToDataY && canSnapTooltipToDataY}
               renderTooltip={renderTooltip}
               renderInPortal={renderTooltipInPortal}
             />
@@ -279,7 +283,6 @@ export default function Example() {
           <button
             onClick={() => {
               setDataMultiplier(5 * Math.random());
-              // setCurrData(currData === data ? halfData : data);
             }}
           >
             Update data
@@ -349,6 +352,7 @@ export default function Example() {
           <label>
             <input
               type="checkbox"
+              disabled={!canSnapTooltipToDataX}
               checked={snapTooltipToDataX}
               onChange={() => setSnapTooltipToDataX(!snapTooltipToDataX)}
             />
@@ -357,6 +361,7 @@ export default function Example() {
           <label>
             <input
               type="checkbox"
+              disabled={!canSnapTooltipToDataY}
               checked={snapTooltipToDataY}
               onChange={() => setSnapTooltipToDataY(!snapTooltipToDataY)}
             />
@@ -505,11 +510,11 @@ export default function Example() {
             />{' '}
             right
           </label>
-          {/* <br />
-        <label>
-          <input type="checkbox" onChange={() => setAutoWidth(!autoWidth)} checked={autoWidth} />{' '}
-          responsive width
-        </label> */}
+          <br />
+          <label>
+            <input type="checkbox" onChange={() => setAutoWidth(!autoWidth)} checked={autoWidth} />{' '}
+            responsive width
+          </label>
         </div>
         <div>
           <strong>series</strong>
diff --git a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/providers/ChartProvider.tsx b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/providers/ChartProvider.tsx
index 08b57dedd..1b2b4969e 100644
--- a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/providers/ChartProvider.tsx
+++ b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/providers/ChartProvider.tsx
@@ -10,6 +10,7 @@ import {
   RegisterData,
   Margin,
   DatumWithKey,
+  ScaleType,
 } from '../../types';
 import ChartContext from '../../context/ChartContext';
 import createScale from '../../createScale';
@@ -69,7 +70,7 @@ export default class ChartProvider<
   }
 
   /** Adds data to the registry and to combined data if it supports events. */
-  registerData: RegisterData = dataToRegister => {
+  registerData: RegisterData<XScaleInput, YScaleInput, Datum> = dataToRegister => {
     this.setState(state => {
       const nextState = {
         ...state,
@@ -156,31 +157,21 @@ export default class ChartProvider<
 
     if (width == null || height == null) return;
 
-    let xDomain = combinedData.map(({ key, datum }) =>
-      dataRegistry[key]?.xAccessor(datum),
-    ) as XScaleInput[];
-
-    let yDomain = combinedData.map(({ key, datum }) =>
-      dataRegistry[key]?.yAccessor(datum),
-    ) as YScaleInput[];
-
-    // apply any updates to the domain from the registry
-    Object.values(dataRegistry).forEach(registry => {
-      if (registry.xDomain) xDomain = registry.xDomain(xDomain);
-      if (registry.yDomain) yDomain = registry.yDomain(yDomain);
-    });
-
-    const xScale = createScale<XScaleInput>({
-      data: xDomain,
+    let xScale = createScale<XScaleInput>({
+      data: combinedData.map(({ key, datum }) =>
+        dataRegistry[key]?.xAccessor(datum),
+      ) as XScaleInput[],
       scaleConfig: xScaleConfig,
       range: [margin.left, width - margin.right],
-    });
+    }) as ScaleType<XScaleInput, number>;
 
-    const yScale = createScale<YScaleInput>({
-      data: yDomain,
+    let yScale = createScale<YScaleInput>({
+      data: combinedData.map(({ key, datum }) =>
+        dataRegistry[key]?.yAccessor(datum),
+      ) as YScaleInput[],
       scaleConfig: yScaleConfig,
       range: [height - margin.bottom, margin.top],
-    });
+    }) as ScaleType<YScaleInput, number>;
 
     const colorScale = scaleOrdinal({
       domain: Object.keys(dataRegistry),
@@ -188,6 +179,13 @@ export default class ChartProvider<
       ...colorScaleConfig,
     });
 
+    // apply any updates to the scales from the registry
+    // @TODO this order currently overrides any changes from x/yScaleConfig
+    Object.values(dataRegistry).forEach(registry => {
+      if (registry.xScale) xScale = registry.xScale(xScale);
+      if (registry.yScale) yScale = registry.yScale(yScale);
+    });
+
     return { xScale, yScale, colorScale };
   };
 
@@ -199,6 +197,7 @@ export default class ChartProvider<
     }
   };
 
+  /**  */
   findNearestData = (event: React.MouseEvent | React.TouchEvent) => {
     const { width, height, margin, xScale, yScale, dataRegistry } = this.state;
 
diff --git a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/series/Stack.tsx b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/series/Stack.tsx
index 3b4f51e26..80dc622e5 100644
--- a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/series/Stack.tsx
+++ b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/components/series/Stack.tsx
@@ -4,7 +4,12 @@ import BarStack from '@vx/shape/lib/shapes/BarStack';
 import BarStackHorizontal from '@vx/shape/lib/shapes/BarStackHorizontal';
 import { BarStack as BarStackType } from '@vx/shape/lib/types';
 import ChartContext from '../../context/ChartContext';
-import { DataRegistry, ChartContext as ChartContextType, NearestDatumArgs } from '../../types';
+import {
+  DataRegistry,
+  ChartContext as ChartContextType,
+  NearestDatumArgs,
+  ScaleType,
+} from '../../types';
 
 import BarSeries from './BarSeries';
 import findNearestDatumY from '../../util/findNearestDatumY';
@@ -13,6 +18,10 @@ import AnimatedBars from './AnimatedBars';
 
 const STACK_ACCESSOR = d => d.stack;
 
+type CombinedData<XScaleInput, YScaleInput> = {
+  [dataKey: string]: XScaleInput | YScaleInput | number;
+} & { stack: XScaleInput | YScaleInput; positiveSum: number; negativeSum: number };
+
 export type GroupProps = {
   horizontal?: boolean;
   children: typeof BarSeries;
@@ -38,7 +47,7 @@ export default function Stack<Datum, XScaleInput, YScaleInput>({
   // override the findNearestDatum logic
   const findNearestDatum = useCallback(
     (args: NearestDatumArgs<Datum, XScaleInput, YScaleInput>) => {
-      if (!stacks.current) return;
+      if (!stacks.current) return null;
 
       const nearestDatum = horizontal
         ? findNearestDatumY<Datum, XScaleInput, YScaleInput>(args)
@@ -84,8 +93,10 @@ export default function Stack<Datum, XScaleInput, YScaleInput>({
   );
 
   // group all child data by stack value, this format is needed by BarStack
-  const combinedData: { [key: string]: number }[] = useMemo(() => {
-    const dataByStackValue = {};
+  const combinedData: CombinedData<XScaleInput, YScaleInput>[] = useMemo(() => {
+    const dataByStackValue: {
+      [stackValue: string]: CombinedData<XScaleInput, YScaleInput>;
+    } = {};
     React.Children.forEach(children, child => {
       const { dataKey, data = [], xAccessor, yAccessor } = child.props;
 
@@ -127,8 +138,17 @@ export default function Stack<Datum, XScaleInput, YScaleInput>({
 
       // only need to update the domain for one of the keys
       if (comprehensiveDomain.length > 0 && dataKeys.indexOf(key) === 0) {
-        dataToRegister[key][horizontal ? 'xDomain' : 'yDomain'] = domain =>
-          extent([...domain, ...comprehensiveDomain], d => d);
+        if (horizontal) {
+          dataToRegister[key].xScale = (scale: ScaleType<XScaleInput, number>) =>
+            scale.domain(
+              extent<number | XScaleInput>([...scale.domain(), ...comprehensiveDomain], d => d),
+            );
+        } else {
+          dataToRegister[key].yScale = (scale: ScaleType<YScaleInput, number>) =>
+            scale.domain(
+              extent<number | YScaleInput>([...scale.domain(), ...comprehensiveDomain], d => d),
+            );
+        }
       }
     });
 
diff --git a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/types/context.ts b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/types/context.ts
index 175b6e870..a038208ee 100644
--- a/packages/vx-demo/src/sandboxes/vx-chart-poc/src/types/context.ts
+++ b/packages/vx-demo/src/sandboxes/vx-chart-poc/src/types/context.ts
@@ -15,10 +15,10 @@ export interface DataRegistry<Datum = unknown, XScaleInput = unknown, YScaleInpu
     yAccessor: (d: Datum) => YScaleInput;
     /** whether the entry supports mouse events. */
     mouseEvents?: boolean;
-    /** Optionally update the domain of the xScale. */
-    xDomain?: (domain: XScaleInput[]) => XScaleInput[];
-    /** Optionally update the domain of the yScale. */
-    yDomain?: (domain: YScaleInput[]) => YScaleInput[];
+    /** Optionally update the xScale. */
+    xScale?: (xScale: ScaleType<XScaleInput, number>) => ScaleType<XScaleInput, number>;
+    /** Optionally update the yScale. */
+    yScale?: (yScale: ScaleType<YScaleInput, number>) => ScaleType<YScaleInput, number>;
     /** Optionally override logic for finding the nearest data point to a mouse event. */
     findNearestDatum?: FindNearestDatum<Datum, XScaleInput, YScaleInput>;
     /** Legend shape */
@@ -26,7 +26,9 @@ export interface DataRegistry<Datum = unknown, XScaleInput = unknown, YScaleInpu
   };
 }
 
-export type RegisterData = (data: DataRegistry) => void;
+export type RegisterData<Datum, XScaleInput, YScaleInput> = (
+  data: DataRegistry<Datum, XScaleInput, YScaleInput>,
+) => void;
 
 export type DatumWithKey<Datum = unknown> = { datum: Datum; key: string; index: number };
 
@@ -57,7 +59,7 @@ export interface ChartContext<
   height: number | null;
   margin: Margin;
   dataRegistry: DataRegistry<Datum, XScaleInput, YScaleInput>;
-  registerData: RegisterData;
+  registerData: RegisterData<Datum, XScaleInput, YScaleInput>;
   unregisterData: (keyOrKeys: string | string[]) => void;
   setChartDimensions: (dims: { width: number; height: number; margin: Margin }) => void;
   findNearestData: (