[charts-pro] Allow exporting a heatmap chart

docs/data/charts/export/ExportChartAsImage.tsx

@@ -73,7 +73,7 @@ function ExportParamsSelector({
 }
 
 export default function ExportChartAsImage() {
-  const apiRef = React.useRef<ChartProApi>(undefined);
+  const apiRef = React.useRef<ChartProApi<'line'>>(undefined);
 
   return (
     <Stack width="100%" gap={2}>

docs/data/charts/export/ExportCompositionNoSnap.tsx

@@ -16,7 +16,7 @@ function CustomChartWrapper({ children }: React.PropsWithChildren) {
 
   return (
     <div
-      ref={chartRootRef}
+      ref={chartRootRef as React.RefObject<HTMLDivElement>}
       style={{ display: 'flex', flexDirection: 'column', alignItems: 'center' }}
     >
       {children}
@@ -25,7 +25,7 @@ function CustomChartWrapper({ children }: React.PropsWithChildren) {
 }
 
 export default function ExportCompositionNoSnap() {
-  const apiRef = React.useRef<ChartProApi>(undefined);
+  const apiRef = React.useRef<ChartProApi<'composition'>>(undefined);
 
   return (
     <Stack width="100%" sx={{ display: 'block' }}>

docs/data/charts/export/PrintChart.js

@@ -1,11 +1,19 @@
 import * as React from 'react';
 import Stack from '@mui/material/Stack';
 import Button from '@mui/material/Button';
+import Select from '@mui/material/Select';
+import InputLabel from '@mui/material/InputLabel';
+import FormControl from '@mui/material/FormControl';
+import MenuItem from '@mui/material/MenuItem';
 import { ScatterChartPro } from '@mui/x-charts-pro/ScatterChartPro';
 
+import { BarChartPro } from '@mui/x-charts-pro/BarChartPro';
+import { LineChartPro } from '@mui/x-charts-pro/LineChartPro';
+import { Heatmap } from '@mui/x-charts-pro/Heatmap';
 import { data } from './randomData';
+import { heatmapData } from './heatmapData';
 
-const series = [
+const scatterSeries = [
   {
     label: 'Series A',
     data: data.map((v) => ({ x: v.x1, y: v.y1, id: v.id })),
@@ -16,19 +24,90 @@ const series = [
   },
 ];
 
+const series = [
+  { label: 'Series A', data: data.map((p) => p.y1) },
+  { label: 'Series B', data: data.map((p) => p.y2) },
+];
+
 export default function PrintChart() {
+  const [chartType, setChartType] = React.useState('scatter');
   const apiRef = React.useRef(undefined);
 
+  const handleChange = (event) => setChartType(event.target.value);
+
   return (
     <Stack width="100%" sx={{ display: 'block' }}>
-      <Button
-        onClick={() => apiRef.current.exportAsPrint()}
-        variant="contained"
+      <Stack
+        width="100%"
+        direction="row"
+        gap={2}
+        justifyContent="center"
         sx={{ mb: 1 }}
       >
-        Print
-      </Button>
-      <ScatterChartPro apiRef={apiRef} height={300} series={series} />
+        <FormControl sx={{ minWidth: 200 }}>
+          <InputLabel id="chart-type-label">Chart Type</InputLabel>
+          <Select
+            labelId="chart-type-label"
+            id="chart-type-select"
+            value={chartType}
+            label="Chart Type"
+            onChange={handleChange}
+          >
+            <MenuItem value="scatter">Scatter</MenuItem>
+            <MenuItem value="line">Line</MenuItem>
+            <MenuItem value="bar">Bar</MenuItem>
+            <MenuItem value="heatmap">Heatmap</MenuItem>
+          </Select>
+        </FormControl>
+        <Button onClick={() => apiRef.current.exportAsPrint()} variant="contained">
+          Print
+        </Button>
+      </Stack>
+      <Chart apiRef={apiRef} type={chartType} />
     </Stack>
   );
 }
+
+function Chart({ apiRef, type }) {
+  switch (type) {
+    case 'scatter':
+      return <ScatterChartPro apiRef={apiRef} height={300} series={scatterSeries} />;
+
+    case 'line':
+      return (
+        <LineChartPro
+          apiRef={apiRef}
+          height={300}
+          xAxis={[{ data: data.map((p) => p.x1).toSorted((a, b) => a - b) }]}
+          series={series}
+        />
+      );
+
+    case 'bar':
+      return (
+        <BarChartPro
+          apiRef={apiRef}
+          height={300}
+          xAxis={[
+            { data: data.map((p) => Math.round(p.x1)).toSorted((a, b) => a - b) },
+          ]}
+          series={series}
+        />
+      );
+
+    case 'heatmap':
+      return (
+        <Heatmap
+          apiRef={apiRef}
+          xAxis={[{ data: [1, 2, 3, 4] }]}
+          yAxis={[{ data: ['A', 'B', 'C', 'D', 'E'] }]}
+          series={[{ data: heatmapData }]}
+          height={300}
+          hideLegend={false}
+        />
+      );
+
+    default:
+      throw new Error(`Unknown chart type: ${type}`);
+  }
+}

docs/data/charts/export/PrintChart.tsx

@@ -1,11 +1,19 @@
 import * as React from 'react';
 import Stack from '@mui/material/Stack';
 import Button from '@mui/material/Button';
+import Select, { SelectChangeEvent } from '@mui/material/Select';
+import InputLabel from '@mui/material/InputLabel';
+import FormControl from '@mui/material/FormControl';
+import MenuItem from '@mui/material/MenuItem';
 import { ScatterChartPro } from '@mui/x-charts-pro/ScatterChartPro';
 import { ChartProApi } from '@mui/x-charts-pro/ChartContainerPro';
+import { BarChartPro } from '@mui/x-charts-pro/BarChartPro';
+import { LineChartPro } from '@mui/x-charts-pro/LineChartPro';
+import { Heatmap } from '@mui/x-charts-pro/Heatmap';
 import { data } from './randomData';
+import { heatmapData } from './heatmapData';
 
-const series = [
+const scatterSeries = [
   {
     label: 'Series A',
     data: data.map((v) => ({ x: v.x1, y: v.y1, id: v.id })),
@@ -15,20 +23,101 @@ const series = [
     data: data.map((v) => ({ x: v.x1, y: v.y2, id: v.id })),
   },
 ];
+const series = [
+  { label: 'Series A', data: data.map((p) => p.y1) },
+  { label: 'Series B', data: data.map((p) => p.y2) },
+];
+
+type ChartType = 'scatter' | 'line' | 'bar' | 'heatmap';
 
 export default function PrintChart() {
+  const [chartType, setChartType] = React.useState<ChartType>('scatter');
   const apiRef = React.useRef<ChartProApi>(undefined);
 
+  const handleChange = (event: SelectChangeEvent) =>
+    setChartType(event.target.value as ChartType);
+
   return (
     <Stack width="100%" sx={{ display: 'block' }}>
-      <Button
-        onClick={() => apiRef.current!.exportAsPrint()}
-        variant="contained"
+      <Stack
+        width="100%"
+        direction="row"
+        gap={2}
+        justifyContent="center"
         sx={{ mb: 1 }}
       >
-        Print
-      </Button>
-      <ScatterChartPro apiRef={apiRef} height={300} series={series} />
+        <FormControl sx={{ minWidth: 200 }}>
+          <InputLabel id="chart-type-label">Chart Type</InputLabel>
+          <Select
+            labelId="chart-type-label"
+            id="chart-type-select"
+            value={chartType}
+            label="Chart Type"
+            onChange={handleChange}
+          >
+            <MenuItem value="scatter">Scatter</MenuItem>
+            <MenuItem value="line">Line</MenuItem>
+            <MenuItem value="bar">Bar</MenuItem>
+            <MenuItem value="heatmap">Heatmap</MenuItem>
+          </Select>
+        </FormControl>
+        <Button onClick={() => apiRef.current!.exportAsPrint()} variant="contained">
+          Print
+        </Button>
+      </Stack>
+      <Chart apiRef={apiRef} type={chartType} />
     </Stack>
   );
 }
+
+function Chart<T extends ChartType = ChartType>({
+  apiRef,
+  type,
+}: {
+  apiRef: React.RefObject<ChartProApi<T> | undefined>;
+  type: T;
+}) {
+  switch (type) {
+    case 'scatter':
+      return (
+        <ScatterChartPro
+          apiRef={apiRef as React.RefObject<ChartProApi<'scatter'> | undefined>}
+          height={300}
+          series={scatterSeries}
+        />
+      );
+    case 'line':
+      return (
+        <LineChartPro
+          apiRef={apiRef as React.RefObject<ChartProApi<'line'> | undefined>}
+          height={300}
+          xAxis={[{ data: data.map((p) => p.x1).toSorted((a, b) => a - b) }]}
+          series={series}
+        />
+      );
+    case 'bar':
+      return (
+        <BarChartPro
+          apiRef={apiRef as React.RefObject<ChartProApi<'bar'> | undefined>}
+          height={300}
+          xAxis={[
+            { data: data.map((p) => Math.round(p.x1)).toSorted((a, b) => a - b) },
+          ]}
+          series={series}
+        />
+      );
+    case 'heatmap':
+      return (
+        <Heatmap
+          apiRef={apiRef as React.RefObject<ChartProApi<'heatmap'> | undefined>}
+          xAxis={[{ data: [1, 2, 3, 4] }]}
+          yAxis={[{ data: ['A', 'B', 'C', 'D', 'E'] }]}
+          series={[{ data: heatmapData }]}
+          height={300}
+          hideLegend={false}
+        />
+      );
+    default:
+      throw new Error(`Unknown chart type: ${type}`);
+  }
+}

docs/data/charts/export/export.md

@@ -8,7 +8,7 @@ components: ScatterChartPro, BarChartPro, LineChartPro
 
 <p class="description">Charts can be printed and exported as PDF.</p>
 
-Export is available on the Pro version of the charts: `<LineChartPro />`, `<BarChartPro />`, `<ScatterChartPro />`.
+Export is available for the following charts: `<LineChartPro />`, `<BarChartPro />`, `<ScatterChartPro />`, `<Heatmap />`.
 
 ## Print/Export as PDF
 

docs/data/charts/export/heatmapData.js

@@ -0,0 +1,22 @@
+export const heatmapData = [
+  [0, 0, 10],
+  [0, 1, 20],
+  [0, 2, 40],
+  [0, 3, 90],
+  [0, 4, 70],
+  [1, 0, 30],
+  [1, 1, 50],
+  [1, 2, 10],
+  [1, 3, 70],
+  [1, 4, 40],
+  [2, 0, 50],
+  [2, 1, 20],
+  [2, 2, 90],
+  [2, 3, 20],
+  [2, 4, 70],
+  [3, 0, 40],
+  [3, 1, 50],
+  [3, 2, 20],
+  [3, 3, 70],
+  [3, 4, 90],
+];

docs/data/charts/export/heatmapData.ts

@@ -0,0 +1,24 @@
+import { HeatmapValueType } from '@mui/x-charts-pro/models';
+
+export const heatmapData: HeatmapValueType[] = [
+  [0, 0, 10],
+  [0, 1, 20],
+  [0, 2, 40],
+  [0, 3, 90],
+  [0, 4, 70],
+  [1, 0, 30],
+  [1, 1, 50],
+  [1, 2, 10],
+  [1, 3, 70],
+  [1, 4, 40],
+  [2, 0, 50],
+  [2, 1, 20],
+  [2, 2, 90],
+  [2, 3, 20],
+  [2, 4, 70],
+  [3, 0, 40],
+  [3, 1, 50],
+  [3, 2, 20],
+  [3, 3, 70],
+  [3, 4, 90],
+];

packages/x-charts-pro/src/ChartContainerPro/ChartContainerPro.tsx

@@ -19,8 +19,6 @@ type ChartContainerProComponent = <
   props: ChartContainerProProps<TSeries, TSignatures> & { ref?: React.ForwardedRef<SVGSVGElement> },
 ) => React.JSX.Element;
 
-export type ChartProApi = NonNullable<NonNullable<ChartContainerProProps['apiRef']>['current']>;
-
 /**
  * It sets up the data providers as well as the `<svg>` for the chart.
  *

packages/x-charts-pro/src/ChartContainerPro/ChartProApi.ts

@@ -0,0 +1,28 @@
+import { ChartAnyPluginSignature, ChartPublicAPI } from '@mui/x-charts/internals';
+import { HeatmapPluginsSignatures } from '../Heatmap/Heatmap.plugins';
+import { LineChartProPluginsSignatures } from '../LineChartPro/LineChartPro.plugins';
+import { ScatterChartProPluginsSignatures } from '../ScatterChartPro/ScatterChartPro.plugins';
+import { BarChartProPluginsSignatures } from '../BarChartPro/BarChartPro.plugins';
+import { AllPluginSignatures, DefaultPluginSignatures } from '../internals/plugins/allPlugins';
+
+type PluginsPerSeriesType = {
+  heatmap: HeatmapPluginsSignatures;
+  line: LineChartProPluginsSignatures;
+  scatter: ScatterChartProPluginsSignatures;
+  bar: BarChartProPluginsSignatures;
+  /* Special value when creating a chart using composition. */
+  composition: DefaultPluginSignatures;
+};
+
+/**
+ * The API of the chart `apiRef` object.
+ * The chart type can be passed as the first generic parameter to narrow down the API to the specific chart type.
+ * @example ChartProApi<'bar'>
+ */
+export type ChartProApi<
+  TSeries extends keyof PluginsPerSeriesType | undefined = undefined,
+  TSignatures extends
+    readonly ChartAnyPluginSignature[] = TSeries extends keyof PluginsPerSeriesType
+    ? PluginsPerSeriesType[TSeries]
+    : AllPluginSignatures,
+> = ChartPublicAPI<TSignatures>;

packages/x-charts-pro/src/ChartContainerPro/index.ts

@@ -1,3 +1,4 @@
 import '../typeOverloads';
 
 export * from './ChartContainerPro';
+export * from './ChartProApi';