diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.ts b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.ts index 100d7bddddeb..9f7d70c11ca7 100644 --- a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.ts +++ b/superset-frontend/temporary_superset_ui/superset-ui/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; diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.ts b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.ts index 0ce297f894e2..00d8244435f8 100644 --- a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.ts +++ b/superset-frontend/temporary_superset_ui/superset-ui/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,6 +100,54 @@ 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']); @@ -106,4 +155,15 @@ describe('CategoricalColorScale', () => { 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'); + }); + }); });