feat: make CategoricalScale compatible with D3 ScaleOrdinal (#357)

* feat: make categorical scale compatible with d3 scaleOrdinal

* feat: make CategoricalColorScale signature compatible with D3 ScaleOrdinal

* test: add unit test

* feat: use scaleOrdinal in implementation

packages/superset-ui-color/src/CategoricalColorScale.ts

@@ -1,16 +1,25 @@
+/* eslint-disable no-dupe-class-members */
 import { ExtensibleFunction } from '@superset-ui/core';
+import { scaleOrdinal, ScaleOrdinal } from 'd3-scale';
 import { ColorsLookup } from './types';
 import stringifyAndTrim from './stringifyAndTrim';
 
-export default class CategoricalColorScale extends ExtensibleFunction {
+// Use type augmentation to correct the fact that
+// an instance of CategoricalScale is also a function
+
+interface CategoricalColorScale {
+  (x: { toString(): string }): string;
+}
+
+class CategoricalColorScale extends ExtensibleFunction {
   colors: string[];
 
+  scale: ScaleOrdinal<{ toString(): string }, string>;
+
   parentForcedColors?: ColorsLookup;
 
   forcedColors: ColorsLookup;
 
-  seen: { [key: string]: number };
-
   /**
    * Constructor
    * @param {*} colors an array of colors
@@ -19,10 +28,12 @@ export default class CategoricalColorScale extends ExtensibleFunction {
    */
   constructor(colors: string[], parentForcedColors?: ColorsLookup) {
     super((value: string) => this.getColor(value));
+
     this.colors = colors;
+    this.scale = scaleOrdinal<{ toString(): string }, string>();
+    this.scale.range(colors);
     this.parentForcedColors = parentForcedColors;
     this.forcedColors = {};
-    this.seen = {};
   }
 
   getColor(value?: string) {
@@ -38,16 +49,7 @@ export default class CategoricalColorScale extends ExtensibleFunction {
       return forcedColor;
     }
 
-    const seenColor = this.seen[cleanedValue];
-    const { length } = this.colors;
-    if (seenColor !== undefined) {
-      return this.colors[seenColor % length];
-    }
-
-    const index = Object.keys(this.seen).length;
-    this.seen[cleanedValue] = index;
-
-    return this.colors[index % length];
+    return this.scale(cleanedValue);
   }
 
   /**
@@ -67,9 +69,8 @@ export default class CategoricalColorScale extends ExtensibleFunction {
    */
   getColorMap() {
     const colorMap: { [key: string]: string } = {};
-    const { length } = this.colors;
-    Object.keys(this.seen).forEach(value => {
-      colorMap[value] = this.colors[this.seen[value] % length];
+    this.scale.domain().forEach(value => {
+      colorMap[value.toString()] = this.scale(value);
     });
 
     return {
@@ -78,4 +79,85 @@ export default class CategoricalColorScale extends ExtensibleFunction {
       ...this.parentForcedColors,
     };
   }
+
+  /**
+   * Returns an exact copy of this scale. Changes to this scale will not affect the returned scale, and vice versa.
+   */
+  copy() {
+    const copy = new CategoricalColorScale(this.scale.range(), this.parentForcedColors);
+    copy.forcedColors = { ...this.forcedColors };
+    copy.domain(this.domain());
+    copy.unknown(this.unknown());
+
+    return copy;
+  }
+
+  /**
+   * Returns the scale's current domain.
+   */
+  domain(): { toString(): string }[];
+
+  /**
+   * Expands the domain to include the specified array of values.
+   */
+  domain(newDomain: { toString(): string }[]): this;
+
+  domain(newDomain?: { toString(): string }[]): unknown {
+    if (typeof newDomain === 'undefined') {
+      return this.scale.domain();
+    }
+
+    this.scale.domain(newDomain);
+    return this;
+  }
+
+  /**
+   * Returns the scale's current range.
+   */
+  range(): string[];
+
+  /**
+   * Sets the range of the ordinal scale to the specified array of values.
+   *
+   * The first element in the domain will be mapped to the first element in range, the second domain value to the second range value, and so on.
+   *
+   * 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.
+   */
+  range(newRange: string[]): this;
+
+  range(newRange?: string[]): unknown {
+    if (typeof newRange === 'undefined') {
+      return this.scale.range();
+    }
+
+    this.colors = newRange;
+    this.scale.range(newRange);
+    return this;
+  }
+
+  /**
+   * Returns the current unknown value, which defaults to "implicit".
+   */
+  unknown(): string | { name: 'implicit' };
+
+  /**
+   * Sets the output value of the scale for unknown input values and returns this scale.
+   * The implicit value enables implicit domain construction. scaleImplicit can be used as a convenience to set the implicit value.
+   *
+   * @param value Unknown value to be used or scaleImplicit to set implicit scale generation.
+   */
+  unknown(value: string | { name: 'implicit' }): this;
+
+  unknown(value?: string | { name: 'implicit' }): unknown {
+    if (typeof value === 'undefined') {
+      return this.scale.unknown();
+    }
+
+    this.scale.unknown(value);
+    return this;
+  }
 }
+
+export default CategoricalColorScale;

packages/superset-ui-color/test/CategoricalColorScale.test.ts

@@ -1,3 +1,4 @@
+import { ScaleOrdinal } from 'd3-scale';
 import CategoricalColorScale from '../src/CategoricalColorScale';
 
 describe('CategoricalColorScale', () => {
@@ -99,11 +100,70 @@ describe('CategoricalColorScale', () => {
       });
     });
   });
+
+  describe('.copy()', () => {
+    it('returns a copy', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      const copy = scale.copy();
+      expect(copy).not.toBe(scale);
+      expect(copy('cat')).toEqual(scale('cat'));
+      expect(copy.domain()).toEqual(scale.domain());
+      expect(copy.range()).toEqual(scale.range());
+      expect(copy.unknown()).toEqual(scale.unknown());
+    });
+  });
+  describe('.domain()', () => {
+    it('when called without argument, returns domain', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      scale.getColor('pig');
+      expect(scale.domain()).toEqual(['pig']);
+    });
+    it('when called with argument, sets domain', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      scale.domain(['dog', 'pig', 'cat']);
+      expect(scale('pig')).toEqual('red');
+    });
+  });
+  describe('.range()', () => {
+    it('when called without argument, returns range', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      expect(scale.range()).toEqual(['blue', 'red', 'green']);
+    });
+    it('when called with argument, sets range', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      scale.range(['pink', 'gray', 'yellow']);
+      expect(scale.range()).toEqual(['pink', 'gray', 'yellow']);
+    });
+  });
+  describe('.unknown()', () => {
+    it('when called without argument, returns output for unknown value', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      scale.unknown('#666');
+      expect(scale.unknown()).toEqual('#666');
+    });
+    it('when called with argument, sets output for unknown value', () => {
+      const scale = new CategoricalColorScale(['blue', 'red', 'green']);
+      scale.unknown('#222');
+      expect(scale.unknown()).toEqual('#222');
+    });
+  });
+
   describe('a CategoricalColorScale instance is also a color function itself', () => {
     it('scale(value) returns color similar to calling scale.getColor(value)', () => {
       const scale = new CategoricalColorScale(['blue', 'red', 'green']);
       expect(scale.getColor('pig')).toBe(scale('pig'));
       expect(scale.getColor('cat')).toBe(scale('cat'));
     });
   });
+
+  describe("is compatible with D3's ScaleOrdinal", () => {
+    it('passes type check', () => {
+      const scale: ScaleOrdinal<{ toString(): string }, string> = new CategoricalColorScale([
+        'blue',
+        'red',
+        'green',
+      ]);
+      expect(scale('pig')).toBe('blue');
+    });
+  });
 });