From 114be9df07f2a7ad6b8fab236f4c378d4a90f556 Mon Sep 17 00:00:00 2001
From: Felipe Jacob Custodio Romero <felipe@neodyme.io>
Date: Mon, 22 Jan 2024 22:27:03 +0100
Subject: [PATCH] [@mantine/charts] Add data formatter for legend and tooltip

---
 .../charts/src/AreaChart/AreaChart.tsx        |  6 +-
 .../@mantine/charts/src/BarChart/BarChart.tsx |  6 +-
 .../charts/src/ChartLegend/ChartLegend.tsx    | 27 ++++---
 .../charts/src/ChartTooltip/ChartTooltip.tsx  | 73 ++++++++++++++-----
 .../charts/src/LineChart/LineChart.tsx        |  6 +-
 5 files changed, 82 insertions(+), 36 deletions(-)

diff --git a/packages/@mantine/charts/src/AreaChart/AreaChart.tsx b/packages/@mantine/charts/src/AreaChart/AreaChart.tsx
index 88d7ddb08f..6ac95ada6c 100644
--- a/packages/@mantine/charts/src/AreaChart/AreaChart.tsx
+++ b/packages/@mantine/charts/src/AreaChart/AreaChart.tsx
@@ -315,13 +315,14 @@ export const AreaChart = factory<AreaChartFactory>((_props, ref) => {
           {withLegend && (
             <Legend
               verticalAlign="top"
-              content={(payload) => (
+              content={({ payload }) => (
                 <ChartLegend
-                  payload={payload.payload}
+                  payload={payload}
                   onHighlight={setHighlightedArea}
                   legendPosition={legendProps?.verticalAlign || 'top'}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={legendProps?.formatter}
                 />
               )}
               height={44}
@@ -379,6 +380,7 @@ export const AreaChart = factory<AreaChartFactory>((_props, ref) => {
                   unit={unit}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={tooltipProps?.formatter}
                 />
               )}
               {...tooltipProps}
diff --git a/packages/@mantine/charts/src/BarChart/BarChart.tsx b/packages/@mantine/charts/src/BarChart/BarChart.tsx
index f7da22e936..d40e98dc19 100644
--- a/packages/@mantine/charts/src/BarChart/BarChart.tsx
+++ b/packages/@mantine/charts/src/BarChart/BarChart.tsx
@@ -223,13 +223,14 @@ export const BarChart = factory<BarChartFactory>((_props, ref) => {
           {withLegend && (
             <Legend
               verticalAlign="top"
-              content={(payload) => (
+              content={({ payload }) => (
                 <ChartLegend
-                  payload={payload.payload}
+                  payload={payload}
                   onHighlight={setHighlightedArea}
                   legendPosition={legendProps?.verticalAlign || 'top'}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={legendProps?.formatter}
                 />
               )}
               height={44}
@@ -288,6 +289,7 @@ export const BarChart = factory<BarChartFactory>((_props, ref) => {
                   unit={unit}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={tooltipProps?.formatter}
                 />
               )}
               {...tooltipProps}
diff --git a/packages/@mantine/charts/src/ChartLegend/ChartLegend.tsx b/packages/@mantine/charts/src/ChartLegend/ChartLegend.tsx
index 26e2fda9e3..741845bd51 100644
--- a/packages/@mantine/charts/src/ChartLegend/ChartLegend.tsx
+++ b/packages/@mantine/charts/src/ChartLegend/ChartLegend.tsx
@@ -1,4 +1,5 @@
 import React from 'react';
+import { Formatter, Payload } from 'recharts/types/component/DefaultLegendContent';
 import {
   Box,
   BoxProps,
@@ -12,7 +13,7 @@ import {
 } from '@mantine/core';
 import classes from './ChartLegend.module.css';
 
-export function getFilteredChartLegendPayload(payload: Record<string, any>[]) {
+export function getFilteredChartLegendPayload(payload: Payload[]) {
   return payload.filter((item) => item.color !== 'none');
 }
 
@@ -22,9 +23,10 @@ export interface ChartLegendProps
   extends BoxProps,
     StylesApiProps<ChartLegendFactory>,
     ElementProps<'div'> {
-  payload: Record<string, any>[] | undefined;
+  payload: Payload[] | undefined;
   onHighlight: (area: string | null) => void;
   legendPosition: 'top' | 'bottom' | 'middle';
+  formatter?: Formatter;
 }
 
 export type ChartLegendFactory = Factory<{
@@ -47,6 +49,7 @@ export const ChartLegend = factory<ChartLegendFactory>((_props, ref) => {
     payload,
     onHighlight,
     legendPosition,
+    formatter,
     ...others
   } = props;
 
@@ -71,16 +74,20 @@ export const ChartLegend = factory<ChartLegendFactory>((_props, ref) => {
     <div
       key={index}
       {...getStyles('legendItem')}
-      onMouseEnter={() => onHighlight(item.dataKey)}
+      onMouseEnter={() => onHighlight(item.value)}
       onMouseLeave={() => onHighlight(null)}
     >
-      <ColorSwatch
-        color={item.color}
-        size={12}
-        {...getStyles('legendItemColor')}
-        withShadow={false}
-      />
-      <p {...getStyles('legendItemName')}>{item.dataKey}</p>
+      {item.color ? (
+        <ColorSwatch
+          color={item.color}
+          size={12}
+          {...getStyles('legendItemColor')}
+          withShadow={false}
+        />
+      ) : undefined}
+      <p {...getStyles('legendItemName')}>
+        {formatter ? formatter(item.value, item, index) : item.value}
+      </p>
     </div>
   ));
 
diff --git a/packages/@mantine/charts/src/ChartTooltip/ChartTooltip.tsx b/packages/@mantine/charts/src/ChartTooltip/ChartTooltip.tsx
index ca35fbca8e..583cfa8669 100644
--- a/packages/@mantine/charts/src/ChartTooltip/ChartTooltip.tsx
+++ b/packages/@mantine/charts/src/ChartTooltip/ChartTooltip.tsx
@@ -1,4 +1,5 @@
 import React from 'react';
+import { Formatter, Payload } from 'recharts/types/component/DefaultTooltipContent';
 import {
   Box,
   BoxProps,
@@ -12,7 +13,9 @@ import {
 } from '@mantine/core';
 import classes from './ChartTooltip.module.css';
 
-export function getFilteredChartTooltipPayload(payload: Record<string, any>[]) {
+type MantinePayload = Payload<string, any> & { fill?: string };
+
+export function getFilteredChartTooltipPayload(payload: MantinePayload[]) {
   return payload.filter((item) => item.fill !== 'none');
 }
 
@@ -34,10 +37,13 @@ export interface ChartTooltipProps
   label?: React.ReactNode;
 
   /** Chart data provided by recharts */
-  payload: Record<string, any>[] | undefined;
+  payload: MantinePayload[] | undefined;
 
   /** Data units, provided by parent component */
   unit?: string;
+
+  /** Function used to format the display content */
+  formatter?: Formatter<string, any>;
 }
 
 export type ChartTooltipFactory = Factory<{
@@ -50,8 +56,19 @@ const defaultProps: Partial<ChartTooltipProps> = {};
 
 export const ChartTooltip = factory<ChartTooltipFactory>((_props, ref) => {
   const props = useProps('ChartTooltip', defaultProps, _props);
-  const { classNames, className, style, styles, unstyled, vars, payload, label, unit, ...others } =
-    props;
+  const {
+    classNames,
+    className,
+    style,
+    styles,
+    unstyled,
+    vars,
+    payload,
+    label,
+    unit,
+    formatter,
+    ...others
+  } = props;
 
   const getStyles = useStyles<ChartTooltipFactory>({
     name: 'ChartTooltip',
@@ -70,23 +87,39 @@ export const ChartTooltip = factory<ChartTooltipFactory>((_props, ref) => {
 
   const filteredPayload = getFilteredChartTooltipPayload(payload);
 
-  const items = filteredPayload.map((item) => (
-    <div key={item.name} {...getStyles('tooltipItem')}>
-      <div {...getStyles('tooltipItemBody')}>
-        <ColorSwatch
-          color={item.color}
-          size={12}
-          {...getStyles('tooltipItemColor')}
-          withShadow={false}
-        />
-        <div {...getStyles('tooltipItemName')}>{item.name}</div>
-      </div>
-      <div {...getStyles('tooltipItemData')}>
-        {item.payload[item.dataKey]}
-        {unit}
+  const items = filteredPayload.map((item, index) => {
+    let finalValue: React.ReactNode | string = item.value;
+    let finalName: React.ReactNode | any = item.name;
+
+    if (formatter !== undefined && item.value !== undefined) {
+      const formatted = formatter(item.value, item.name, item, index, payload);
+      if (Array.isArray(formatted)) {
+        [finalValue, finalName] = formatted;
+      } else {
+        finalValue = formatted;
+      }
+    }
+
+    return (
+      <div key={index} {...getStyles('tooltipItem')}>
+        <div {...getStyles('tooltipItemBody')}>
+          {item.color ? (
+            <ColorSwatch
+              color={item.color}
+              size={12}
+              {...getStyles('tooltipItemColor')}
+              withShadow={false}
+            />
+          ) : undefined}
+          <div {...getStyles('tooltipItemName')}>{finalName}</div>
+        </div>
+        <div {...getStyles('tooltipItemData')}>
+          {finalValue}
+          {unit}
+        </div>
       </div>
-    </div>
-  ));
+    );
+  });
 
   return (
     <Box {...getStyles('tooltip')} ref={ref} {...others}>
diff --git a/packages/@mantine/charts/src/LineChart/LineChart.tsx b/packages/@mantine/charts/src/LineChart/LineChart.tsx
index 82c8966ff5..611dafbea1 100644
--- a/packages/@mantine/charts/src/LineChart/LineChart.tsx
+++ b/packages/@mantine/charts/src/LineChart/LineChart.tsx
@@ -248,13 +248,14 @@ export const LineChart = factory<LineChartFactory>((_props, ref) => {
           {withLegend && (
             <Legend
               verticalAlign="top"
-              content={(payload) => (
+              content={({ payload }) => (
                 <ChartLegend
-                  payload={payload.payload}
+                  payload={payload}
                   onHighlight={setHighlightedArea}
                   legendPosition={legendProps?.verticalAlign || 'top'}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={legendProps?.formatter}
                 />
               )}
               height={44}
@@ -311,6 +312,7 @@ export const LineChart = factory<LineChartFactory>((_props, ref) => {
                   unit={unit}
                   classNames={resolvedClassNames}
                   styles={resolvedStyles}
+                  formatter={tooltipProps?.formatter}
                 />
               )}
               {...tooltipProps}