Make CategoricalColorScale instance a function and remove .toFunction()

packages/generator-superset/package.json

@@ -33,5 +33,8 @@
     "yeoman-generator": "^3.1.1",
     "yosay": "^2.0.2"
   },
-  "license": "Apache-2.0"
+  "license": "Apache-2.0",
+  "publishConfig": {
+    "access": "public"
+  }
 }

packages/superset-ui-color/src/CategoricalColorNamespace.js

@@ -1,5 +1,6 @@
 import CategoricalColorScale from './CategoricalColorScale';
 import getCategoricalSchemeRegistry from './CategoricalSchemeRegistrySingleton';
+import cleanValue from './cleanValue';
 
 export default class CategoricalColorNamespace {
   constructor(name) {
@@ -31,7 +32,7 @@ export default class CategoricalColorNamespace {
    * @param {*} forcedColor color
    */
   setColor(value, forcedColor) {
-    this.forcedItems[value] = forcedColor;
+    this.forcedItems[cleanValue(value)] = forcedColor;
 
     return this;
   }

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

@@ -1,23 +1,19 @@
-export function cleanValue(value) {
-  // for superset series that should have the same color
-  return String(value)
-    .trim()
-    .toLowerCase();
-}
+import { ExtensibleFunction } from '@superset-ui/core';
+import cleanValue from './cleanValue';
 
-export default class CategoricalColorScale {
+export default class CategoricalColorScale extends ExtensibleFunction {
   /**
    * Constructor
    * @param {*} colors an array of colors
    * @param {*} parentForcedColors optional parameter that comes from parent
    * (usually CategoricalColorNamespace) and supersede this.forcedColors
    */
   constructor(colors, parentForcedColors) {
+    super((...args) => this.getColor(...args));
     this.colors = colors;
     this.parentForcedColors = parentForcedColors;
     this.forcedColors = {};
     this.seen = {};
-    this.fn = value => this.getColor(value);
   }
 
   getColor(value) {
@@ -51,12 +47,8 @@ export default class CategoricalColorScale {
    * @param {*} forcedColor forcedColor
    */
   setColor(value, forcedColor) {
-    this.forcedColors[value] = forcedColor;
+    this.forcedColors[cleanValue(value)] = forcedColor;
 
     return this;
   }
-
-  toFunction() {
-    return this.fn;
-  }
 }

packages/superset-ui-color/src/cleanValue.js

@@ -0,0 +1,7 @@
+/**
+ * Ensure value is a string
+ * @param {any} value
+ */
+export default function cleanValue(value) {
+  return String(value).trim();
+}

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

@@ -83,12 +83,11 @@ describe('CategoricalColorScale', () => {
       expect(scale).toBe(output);
     });
   });
-  describe('.toFunction()', () => {
-    it('returns a function that wraps getColor', () => {
+  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']);
-      const colorFn = scale.toFunction();
-      expect(scale.getColor('pig')).toBe(colorFn('pig'));
-      expect(scale.getColor('cat')).toBe(colorFn('cat'));
+      expect(scale.getColor('pig')).toBe(scale('pig'));
+      expect(scale.getColor('cat')).toBe(scale('cat'));
     });
   });
 });