diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorNamespace.js b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorNamespace.js
index 218085867cfe..eb52ddfa6824 100644
--- a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorNamespace.js
+++ b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorNamespace.js
@@ -1,5 +1,6 @@
 import CategoricalColorScale from './CategoricalColorScale';
 import getCategoricalSchemeRegistry from './CategoricalSchemeRegistrySingleton';
+import stringifyAndTrim from './stringifyAndTrim';
 
 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[stringifyAndTrim(value)] = forcedColor;
 
     return this;
   }
diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.js b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.js
index ceef5c365ed6..4d37a2da3619 100644
--- a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.js
+++ b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/CategoricalColorScale.js
@@ -1,11 +1,7 @@
-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 stringifyAndTrim from './stringifyAndTrim';
 
-export default class CategoricalColorScale {
+export default class CategoricalColorScale extends ExtensibleFunction {
   /**
    * Constructor
    * @param {*} colors an array of colors
@@ -13,15 +9,15 @@ export default class CategoricalColorScale {
    * (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) {
-    const cleanedValue = cleanValue(value);
+    const cleanedValue = stringifyAndTrim(value);
 
     const parentColor = this.parentForcedColors && this.parentForcedColors[cleanedValue];
     if (parentColor) {
@@ -51,12 +47,8 @@ export default class CategoricalColorScale {
    * @param {*} forcedColor forcedColor
    */
   setColor(value, forcedColor) {
-    this.forcedColors[value] = forcedColor;
+    this.forcedColors[stringifyAndTrim(value)] = forcedColor;
 
     return this;
   }
-
-  toFunction() {
-    return this.fn;
-  }
 }
diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/stringifyAndTrim.js b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/stringifyAndTrim.js
new file mode 100644
index 000000000000..06e4d53b93b1
--- /dev/null
+++ b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/src/stringifyAndTrim.js
@@ -0,0 +1,7 @@
+/**
+ * Ensure value is a string
+ * @param {any} value
+ */
+export default function stringifyAndTrim(value) {
+  return String(value).trim();
+}
diff --git a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.js b/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.js
index 13d238654b53..9581fdcb6305 100644
--- a/superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-color/test/CategoricalColorScale.test.js
+++ b/superset-frontend/temporary_superset_ui/superset-ui/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'));
     });
   });
 });