From c8fa44e9e904160de705cd643d1df092815348b1 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Wed, 12 Apr 2023 14:27:09 +0300
Subject: [PATCH] feat(dashboard): make color indices referable (#23657)

---
 .../src/color/CategoricalColorScale.ts        | 20 ++++++++++++++-----
 .../superset-ui-core/src/color/types.ts       |  4 ++++
 .../test/color/CategoricalColorScale.test.ts  | 15 ++++++++++++++
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/superset-frontend/packages/superset-ui-core/src/color/CategoricalColorScale.ts b/superset-frontend/packages/superset-ui-core/src/color/CategoricalColorScale.ts
index 5f29ad4775..3fc093d559 100644
--- a/superset-frontend/packages/superset-ui-core/src/color/CategoricalColorScale.ts
+++ b/superset-frontend/packages/superset-ui-core/src/color/CategoricalColorScale.ts
@@ -20,7 +20,7 @@
 /* eslint-disable no-dupe-class-members */
 import { scaleOrdinal, ScaleOrdinal } from 'd3-scale';
 import { ExtensibleFunction } from '../models';
-import { ColorsLookup } from './types';
+import { ColorsInitLookup, ColorsLookup } from './types';
 import stringifyAndTrim from './stringifyAndTrim';
 import getSharedLabelColor from './SharedLabelColorSingleton';
 import { getAnalogousColors } from './utils';
@@ -39,7 +39,7 @@ class CategoricalColorScale extends ExtensibleFunction {
 
   scale: ScaleOrdinal<{ toString(): string }, string>;
 
-  parentForcedColors?: ColorsLookup;
+  parentForcedColors: ColorsLookup;
 
   forcedColors: ColorsLookup;
 
@@ -51,14 +51,24 @@ class CategoricalColorScale extends ExtensibleFunction {
    * @param {*} parentForcedColors optional parameter that comes from parent
    * (usually CategoricalColorNamespace) and supersede this.forcedColors
    */
-  constructor(colors: string[], parentForcedColors: ColorsLookup = {}) {
+  constructor(colors: string[], parentForcedColors: ColorsInitLookup = {}) {
     super((value: string, sliceId?: number) => this.getColor(value, sliceId));
 
     this.originColors = colors;
     this.colors = colors;
     this.scale = scaleOrdinal<{ toString(): string }, string>();
     this.scale.range(colors);
-    this.parentForcedColors = parentForcedColors;
+
+    // reserve fixed colors in parent map based on their index in the scale
+    Object.entries(parentForcedColors).forEach(([key, value]) => {
+      if (typeof value === 'number') {
+        // eslint-disable-next-line no-param-reassign
+        parentForcedColors[key] = colors[value % colors.length];
+      }
+    });
+
+    // all indexes have been replaced by a fixed color
+    this.parentForcedColors = parentForcedColors as ColorsLookup;
     this.forcedColors = {};
     this.multiple = 0;
   }
@@ -165,7 +175,7 @@ class CategoricalColorScale extends ExtensibleFunction {
    *
    * If there are fewer elements in the range than in the domain, the scale will reuse values from the start of the range.
    *
-   * @param range Array of range values.
+   * @param newRange Array of range values.
    */
   range(newRange: string[]): this;
 
diff --git a/superset-frontend/packages/superset-ui-core/src/color/types.ts b/superset-frontend/packages/superset-ui-core/src/color/types.ts
index 551c249979..2e4b528f9b 100644
--- a/superset-frontend/packages/superset-ui-core/src/color/types.ts
+++ b/superset-frontend/packages/superset-ui-core/src/color/types.ts
@@ -17,6 +17,10 @@
  * under the License.
  */
 
+export interface ColorsInitLookup {
+  [key: string]: string | number;
+}
+
 export interface ColorsLookup {
   [key: string]: string;
 }
diff --git a/superset-frontend/packages/superset-ui-core/test/color/CategoricalColorScale.test.ts b/superset-frontend/packages/superset-ui-core/test/color/CategoricalColorScale.test.ts
index 9e83aaba9a..df79e778b6 100644
--- a/superset-frontend/packages/superset-ui-core/test/color/CategoricalColorScale.test.ts
+++ b/superset-frontend/packages/superset-ui-core/test/color/CategoricalColorScale.test.ts
@@ -39,7 +39,22 @@ describe('CategoricalColorScale', () => {
       expect(scale).toBeInstanceOf(CategoricalColorScale);
       expect(scale.parentForcedColors).toBe(parentForcedColors);
     });
+
+    it('can refer to colors based on their index', () => {
+      const parentForcedColors = { pig: 1, horse: 5 };
+      const scale = new CategoricalColorScale(
+        ['blue', 'red', 'green'],
+        parentForcedColors,
+      );
+      expect(scale.getColor('pig')).toEqual('red');
+      expect(parentForcedColors.pig).toEqual('red');
+
+      // can loop around the scale
+      expect(scale.getColor('horse')).toEqual('green');
+      expect(parentForcedColors.horse).toEqual('green');
+    });
   });
+
   describe('.getColor(value)', () => {
     it('returns same color for same value', () => {
       const scale = new CategoricalColorScale(['blue', 'red', 'green']);