From 1f8982778a2ae604ee17aaf28f2f0f9b42877e3e Mon Sep 17 00:00:00 2001
From: Ross Gray <ross@posthog.com>
Date: Wed, 30 Oct 2024 10:56:54 +0000
Subject: [PATCH] Add color picker to display settings

---
 frontend/src/lib/colors.ts                    |  5 +++++
 .../Components/Charts/LineGraph.tsx           |  2 +-
 .../Components/ColorPickerButton.tsx          |  8 +++----
 .../Components/SeriesTab.tsx                  | 22 ++++++++++++++-----
 .../Components/ySeriesLogic.ts                |  3 +++
 frontend/src/queries/schema.ts                |  1 +
 6 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/frontend/src/lib/colors.ts b/frontend/src/lib/colors.ts
index cc9530b5bc524a..9646b6dbc953d4 100644
--- a/frontend/src/lib/colors.ts
+++ b/frontend/src/lib/colors.ts
@@ -60,6 +60,11 @@ export function getSeriesColor(index: number = 0): string {
     return getColorVar(`data-${dataColorVars[adjustedIndex]}`)
 }
 
+/** Returns all color options for series */
+export function getSeriesColorPalette(): string[] {
+    return dataColorVars.map((colorVar) => getColorVar(`data-${colorVar}`))
+}
+
 /** Return the background color for the given series index. */
 export function getSeriesBackgroundColor(index: number): string {
     return `${getSeriesColor(index)}30`
diff --git a/frontend/src/queries/nodes/DataVisualization/Components/Charts/LineGraph.tsx b/frontend/src/queries/nodes/DataVisualization/Components/Charts/LineGraph.tsx
index 2cf06e69171445..06b79cd0e1ce37 100644
--- a/frontend/src/queries/nodes/DataVisualization/Components/Charts/LineGraph.tsx
+++ b/frontend/src/queries/nodes/DataVisualization/Components/Charts/LineGraph.tsx
@@ -114,7 +114,7 @@ export const LineGraph = (): JSX.Element => {
         const data: ChartData = {
             labels: xData.data,
             datasets: yData.map(({ data, settings }, index) => {
-                const color = getSeriesColor(index)
+                const color = settings?.display?.color ?? getSeriesColor(index)
                 const backgroundColor = isAreaChart ? hexToRGBA(color, 0.5) : color
 
                 const graphType = getGraphType(visualizationType, settings)
diff --git a/frontend/src/queries/nodes/DataVisualization/Components/ColorPickerButton.tsx b/frontend/src/queries/nodes/DataVisualization/Components/ColorPickerButton.tsx
index 1d44ac883c5ee8..a9fdac3aedf91e 100644
--- a/frontend/src/queries/nodes/DataVisualization/Components/ColorPickerButton.tsx
+++ b/frontend/src/queries/nodes/DataVisualization/Components/ColorPickerButton.tsx
@@ -23,9 +23,11 @@ const DEFAULT_PICKER_COLORS = [
 export const ColorPickerButton = ({
     color,
     onColorSelect: propOnColorSelect,
+    colorChoices = DEFAULT_PICKER_COLORS,
 }: {
     color: string
     onColorSelect?: (color: string) => void
+    colorChoices?: string[]
 }): JSX.Element => {
     const [pickerOpen, setPickerOpen] = useState(false)
     const { isDarkModeOn } = useValues(themeLogic)
@@ -35,14 +37,12 @@ export const ColorPickerButton = ({
             propOnColorSelect(colorResult.hex)
         }
 
-        if (DEFAULT_PICKER_COLORS.includes(colorResult.hex)) {
+        if (colorChoices.includes(colorResult.hex)) {
             setPickerOpen(false)
         }
     }
 
-    const colors = isDarkModeOn
-        ? DEFAULT_PICKER_COLORS.map((n) => RGBToHex(lightenDarkenColor(n, -30)))
-        : DEFAULT_PICKER_COLORS
+    const colors = isDarkModeOn ? colorChoices.map((n) => RGBToHex(lightenDarkenColor(n, -30))) : colorChoices
 
     return (
         <Popover
diff --git a/frontend/src/queries/nodes/DataVisualization/Components/SeriesTab.tsx b/frontend/src/queries/nodes/DataVisualization/Components/SeriesTab.tsx
index bb16cb2cd613b7..b471596a2dad95 100644
--- a/frontend/src/queries/nodes/DataVisualization/Components/SeriesTab.tsx
+++ b/frontend/src/queries/nodes/DataVisualization/Components/SeriesTab.tsx
@@ -12,7 +12,7 @@ import {
 } from '@posthog/lemon-ui'
 import { useActions, useValues } from 'kea'
 import { Form } from 'kea-forms'
-import { getSeriesColor } from 'lib/colors'
+import { getSeriesColor, getSeriesColorPalette } from 'lib/colors'
 import { SeriesGlyph } from 'lib/components/SeriesGlyph'
 import { LemonField } from 'lib/lemon-ui/LemonField'
 import { hexToRGBA, lightenDarkenColor, RGBToRGBA } from 'lib/utils'
@@ -20,6 +20,7 @@ import { hexToRGBA, lightenDarkenColor, RGBToRGBA } from 'lib/utils'
 import { themeLogic } from '~/layout/navigation-3000/themeLogic'
 
 import { AxisSeries, dataVisualizationLogic } from '../dataVisualizationLogic'
+import { ColorPickerButton } from './ColorPickerButton'
 import { ySeriesLogic, YSeriesLogicProps, YSeriesSettingsTab } from './ySeriesLogic'
 
 export const SeriesTab = (): JSX.Element => {
@@ -98,7 +99,7 @@ const YSeries = ({ series, index }: { series: AxisSeries<number>; index: number
     const { setSettingsOpen, submitFormatting, submitDisplay, setSettingsTab } = useActions(seriesLogic)
 
     const { isDarkModeOn } = useValues(themeLogic)
-    const seriesColor = getSeriesColor(index)
+    const seriesColor = series.settings?.display?.color ?? getSeriesColor(index)
 
     const columnsInOptions = showTableSettings ? columns : numericalColumns
     const options = columnsInOptions.map(({ name, type }) => ({
@@ -219,9 +220,20 @@ const YSeriesDisplayTab = ({ ySeriesLogicProps }: { ySeriesLogicProps: YSeriesLo
 
     return (
         <Form logic={ySeriesLogic} props={ySeriesLogicProps} formKey="display" className="space-y-4">
-            <LemonField name="label" label="Label">
-                <LemonInput />
-            </LemonField>
+            <div className="flex gap-3">
+                <LemonField name="color" label="Color">
+                    {({ value, onChange }) => (
+                        <ColorPickerButton
+                            color={value}
+                            onColorSelect={onChange}
+                            colorChoices={getSeriesColorPalette()}
+                        />
+                    )}
+                </LemonField>
+                <LemonField name="label" label="Label">
+                    <LemonInput />
+                </LemonField>
+            </div>
             {!showTableSettings && (
                 <>
                     <LemonField name="trendLine" label="Trend line">
diff --git a/frontend/src/queries/nodes/DataVisualization/Components/ySeriesLogic.ts b/frontend/src/queries/nodes/DataVisualization/Components/ySeriesLogic.ts
index d4f992add18598..e938965aad3c1e 100644
--- a/frontend/src/queries/nodes/DataVisualization/Components/ySeriesLogic.ts
+++ b/frontend/src/queries/nodes/DataVisualization/Components/ySeriesLogic.ts
@@ -1,5 +1,6 @@
 import { actions, connect, kea, key, path, props, reducers, selectors } from 'kea'
 import { forms } from 'kea-forms'
+import { getSeriesColor } from 'lib/colors'
 
 import {
     AxisSeries,
@@ -75,6 +76,7 @@ export const ySeriesLogic = kea<ySeriesLogicType>([
         },
         display: {
             defaults: {
+                color: props.series?.settings?.display?.color ?? getSeriesColor(props.seriesIndex),
                 label: props.series?.settings?.display?.label ?? '',
                 trendLine: props.series?.settings?.display?.trendLine ?? false,
                 yAxisPosition: props.series?.settings?.display?.yAxisPosition ?? 'left',
@@ -83,6 +85,7 @@ export const ySeriesLogic = kea<ySeriesLogicType>([
             submit: async (display) => {
                 actions.updateSeriesIndex(props.seriesIndex, props.series.column.name, {
                     display: {
+                        color: display.color,
                         label: display.label,
                         trendLine: display.trendLine,
                         yAxisPosition: display.yAxisPosition,
diff --git a/frontend/src/queries/schema.ts b/frontend/src/queries/schema.ts
index 273605a42f6d7d..b7a098f5f440b0 100644
--- a/frontend/src/queries/schema.ts
+++ b/frontend/src/queries/schema.ts
@@ -668,6 +668,7 @@ export interface ChartSettingsFormatting {
 }
 
 export interface ChartSettingsDisplay {
+    color?: string
     label?: string
     trendLine?: boolean
     yAxisPosition?: 'left' | 'right'