feat: add function for getting legend information (apache#236)

* feat: add .getLegendInformation()

* fix: lint

* test: add unit test

* feat: revise how legend information is computed

* fix: address comments

superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-encodable/src/encoders/ChannelEncoder.ts

@@ -3,7 +3,10 @@ import { ChannelType, ChannelInput } from '../types/Channel';
 import { PlainObject, Dataset } from '../types/Data';
 import { ChannelDef } from '../types/ChannelDef';
 import createGetterFromChannelDef, { Getter } from '../parsers/createGetterFromChannelDef';
-import completeChannelDef, { CompleteChannelDef } from '../fillers/completeChannelDef';
+import completeChannelDef, {
+  CompleteChannelDef,
+  CompleteValueDef,
+} from '../fillers/completeChannelDef';
 import createFormatterFromChannelDef from '../parsers/format/createFormatterFromChannelDef';
 import createScaleFromScaleConfig from '../parsers/scale/createScaleFromScaleConfig';
 import identity from '../utils/identity';
@@ -44,7 +47,14 @@ export default class ChannelEncoder<Def extends ChannelDef<Output>, Output exten
     this.formatValue = createFormatterFromChannelDef(this.definition);
 
     const scale = this.definition.scale && createScaleFromScaleConfig(this.definition.scale);
-    this.encodeValue = scale === false ? identity : (value: ChannelInput) => scale(value);
+    if (scale === false) {
+      this.encodeValue =
+        'value' in this.definition
+          ? () => (this.definition as CompleteValueDef<Output>).value
+          : identity;
+    } else {
+      this.encodeValue = (value: ChannelInput) => scale(value);
+    }
     this.scale = scale;
   }
 
@@ -80,7 +90,7 @@ export default class ChannelEncoder<Def extends ChannelDef<Output>, Output exten
 
     const { type } = this.definition;
     if (type === 'nominal' || type === 'ordinal') {
-      return Array.from(new Set(data.map(d => this.getValueFromDatum(d)))) as string[];
+      return Array.from(new Set(data.map(d => this.getValueFromDatum(d)))) as ChannelInput[];
     } else if (type === 'quantitative') {
       const extent = d3Extent(data, d => this.getValueFromDatum<number>(d));
 

superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-encodable/src/encoders/Encoder.ts

@@ -1,23 +1,27 @@
 import { flatMap } from 'lodash';
 import { ChannelDef, TypedFieldDef } from '../types/ChannelDef';
 import { MayBeArray } from '../types/Base';
-import { isFieldDef } from '../typeGuards/ChannelDef';
+import { isTypedFieldDef, isValueDef } from '../typeGuards/ChannelDef';
 import { isNotArray } from '../typeGuards/Base';
 import ChannelEncoder from './ChannelEncoder';
 import {
   EncodingConfig,
   DeriveEncoding,
   DeriveChannelTypes,
   DeriveChannelEncoders,
+  DeriveSingleChannelEncoder,
 } from '../types/Encoding';
+import { Dataset } from '../types/Data';
+import { Value } from '../types/VegaLite';
+import { ChannelInput } from '../types/Channel';
 
 export default class Encoder<Config extends EncodingConfig> {
   readonly encoding: DeriveEncoding<Config>;
   readonly channelTypes: DeriveChannelTypes<Config>;
   readonly channels: DeriveChannelEncoders<Config>;
 
   readonly legends: {
-    [key: string]: (keyof Config)[];
+    [key: string]: DeriveSingleChannelEncoder<Config>[];
   };
 
   constructor({
@@ -64,13 +68,12 @@ export default class Encoder<Config extends EncodingConfig> {
     channelNames
       .map(name => this.channels[name])
       .forEach(c => {
-        if (isNotArray(c) && c.hasLegend() && isFieldDef(c.definition)) {
-          const name = c.name as keyof Config;
+        if (isNotArray(c) && c.hasLegend() && isTypedFieldDef(c.definition)) {
           const { field } = c.definition;
           if (this.legends[field]) {
-            this.legends[field].push(name);
+            this.legends[field].push(c);
           } else {
-            this.legends[field] = [name];
+            this.legends[field] = [c];
           }
         }
       });
@@ -92,6 +95,57 @@ export default class Encoder<Config extends EncodingConfig> {
     return Array.from(new Set(fields));
   }
 
+  private createLegendItemsFactory(field: string) {
+    const channelEncoders = flatMap(
+      this.getChannelEncoders().filter(e => isNotArray(e) && isValueDef(e.definition)),
+    ).concat(this.legends[field]);
+
+    return (domain: ChannelInput[]) =>
+      domain.map((input: ChannelInput) => ({
+        input,
+        output: channelEncoders.reduce(
+          (prev: Partial<{ [k in keyof Config]: Config[k]['1'] }>, curr) => {
+            const map = prev;
+            map[curr.name as keyof Config] = curr.encodeValue(input) as Value;
+
+            return map;
+          },
+          {},
+        ),
+      }));
+  }
+
+  getLegendInformation(data: Dataset = []) {
+    return (
+      Object.keys(this.legends)
+        // for each field that was encoded
+        .map((field: string) => {
+          // get all the channels that use this field
+          const channelEncoders = this.legends[field];
+          const firstEncoder = channelEncoders[0];
+          const definition = firstEncoder.definition as TypedFieldDef;
+          const createLegendItems = this.createLegendItemsFactory(field);
+
+          if (definition.type === 'nominal') {
+            return {
+              channelEncoders,
+              createLegendItems,
+              field,
+              items: createLegendItems(firstEncoder.getDomain(data)),
+              type: definition.type,
+            };
+          }
+
+          return {
+            channelEncoders,
+            createLegendItems,
+            field,
+            type: definition.type,
+          };
+        })
+    );
+  }
+
   hasLegend() {
     return Object.keys(this.legends).length > 0;
   }

superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-encodable/src/types/Encoding.ts

@@ -11,7 +11,9 @@ export type DeriveChannelTypes<Config extends EncodingConfig> = {
 };
 
 export type DeriveChannelOutputs<Config extends EncodingConfig> = {
-  readonly [k in keyof Config]: Config[k]['1'];
+  readonly [k in keyof Config]: Config[k]['2'] extends 'multiple'
+    ? Config[k]['1'][]
+    : Config[k]['1'];
 };
 
 export type DeriveEncoding<Config extends EncodingConfig> = {
@@ -25,3 +27,8 @@ export type DeriveChannelEncoders<Config extends EncodingConfig> = {
     ? ChannelEncoder<ChannelTypeToDefMap<Config[k]['1']>[Config[k]['0']]>[]
     : ChannelEncoder<ChannelTypeToDefMap<Config[k]['1']>[Config[k]['0']]>;
 };
+
+export type DeriveSingleChannelEncoder<
+  Config extends EncodingConfig,
+  k extends keyof Config = keyof Config
+> = ChannelEncoder<ChannelTypeToDefMap<Config[k]['1']>[Config[k]['0']]>;

superset-frontend/temporary_superset_ui/superset-ui/packages/superset-ui-encodable/test/encoders/Encoder.test.ts

@@ -1,25 +1,36 @@
 import createEncoderFactory from '../../src/encoders/createEncoderFactory';
 
+function stripFunction(legendInfo) {
+  return legendInfo.map(legendGroup => {
+    const { createLegendItems, channelEncoders, ...rest } = legendGroup;
+
+    return { ...rest };
+  });
+}
+
 describe('Encoder', () => {
   const factory = createEncoderFactory<{
     x: ['X', number];
     y: ['Y', number];
     color: ['Color', string];
+    radius: ['Numeric', number];
     shape: ['Category', string];
     tooltip: ['Text', string, 'multiple'];
   }>({
     channelTypes: {
       x: 'X',
       y: 'Y',
       color: 'Color',
+      radius: 'Numeric',
       shape: 'Category',
       tooltip: 'Text',
     },
     defaultEncoding: {
       x: { type: 'quantitative', field: 'speed' },
       y: { type: 'quantitative', field: 'price' },
       color: { type: 'nominal', field: 'brand' },
-      shape: { type: 'nominal', field: 'brand' },
+      radius: { value: 5 },
+      shape: { value: 'circle' },
       tooltip: [{ field: 'make' }, { field: 'model' }],
     },
   });
@@ -33,19 +44,157 @@ describe('Encoder', () => {
   });
   describe('.getChannelNames()', () => {
     it('returns an array of channel names', () => {
-      expect(encoder.getChannelNames()).toEqual(['x', 'y', 'color', 'shape', 'tooltip']);
+      expect(encoder.getChannelNames()).toEqual(['x', 'y', 'color', 'radius', 'shape', 'tooltip']);
     });
   });
   describe('.getChannelEncoders()', () => {
     it('returns an array of channel encoders', () => {
-      expect(encoder.getChannelEncoders()).toHaveLength(5);
+      expect(encoder.getChannelEncoders()).toHaveLength(6);
     });
   });
   describe('.getGroupBys()', () => {
     it('returns an array of groupby fields', () => {
       expect(encoder.getGroupBys()).toEqual(['brand', 'make', 'model']);
     });
   });
+  describe('.getLegendInformation()', () => {
+    it('returns information for each field', () => {
+      const legendInfo = factory
+        .create({
+          color: { type: 'nominal', field: 'brand', scale: { range: ['red', 'green', 'blue'] } },
+          shape: { type: 'nominal', field: 'brand', scale: { range: ['circle', 'diamond'] } },
+        })
+        .getLegendInformation([{ brand: 'Gucci' }, { brand: 'Prada' }]);
+
+      expect(stripFunction(legendInfo)).toEqual([
+        {
+          field: 'brand',
+          type: 'nominal',
+          items: [
+            {
+              input: 'Gucci',
+              output: {
+                color: 'red',
+                radius: 5,
+                shape: 'circle',
+              },
+            },
+            {
+              input: 'Prada',
+              output: {
+                color: 'green',
+                radius: 5,
+                shape: 'diamond',
+              },
+            },
+          ],
+        },
+      ]);
+    });
+    it('ignore channels that are ValueDef', () => {
+      const legendInfo = factory
+        .create({
+          color: { type: 'nominal', field: 'brand', scale: { range: ['red', 'green', 'blue'] } },
+        })
+        .getLegendInformation([{ brand: 'Gucci' }, { brand: 'Prada' }]);
+
+      expect(stripFunction(legendInfo)).toEqual([
+        {
+          field: 'brand',
+          type: 'nominal',
+          items: [
+            {
+              input: 'Gucci',
+              output: {
+                color: 'red',
+                radius: 5,
+                shape: 'circle',
+              },
+            },
+            {
+              input: 'Prada',
+              output: {
+                color: 'green',
+                radius: 5,
+                shape: 'circle',
+              },
+            },
+          ],
+        },
+      ]);
+    });
+    it('for non-nominal fields, does not return items', () => {
+      const legendInfo = factory
+        .create({
+          color: {
+            type: 'quantitative',
+            field: 'price',
+            scale: { domain: [0, 20], range: ['#fff', '#f00'] },
+          },
+        })
+        .getLegendInformation();
+
+      expect(stripFunction(legendInfo)).toEqual([
+        {
+          field: 'price',
+          type: 'quantitative',
+        },
+      ]);
+    });
+    it('for non-nominal fields, can use createLegendItems function', () => {
+      const legendInfo = factory
+        .create({
+          color: {
+            type: 'quantitative',
+            field: 'price',
+            scale: { domain: [0, 20], range: ['#fff', '#f00'] },
+          },
+          radius: {
+            type: 'quantitative',
+            field: 'price',
+            scale: { domain: [0, 20], range: [0, 10] },
+          },
+        })
+        .getLegendInformation();
+
+      expect(legendInfo[0].createLegendItems([0, 10, 20])).toEqual([
+        {
+          input: 0,
+          output: {
+            color: 'rgb(255, 255, 255)',
+            radius: 0,
+            shape: 'circle',
+          },
+        },
+        {
+          input: 10,
+          output: {
+            color: 'rgb(255, 128, 128)',
+            radius: 5,
+            shape: 'circle',
+          },
+        },
+        {
+          input: 20,
+          output: {
+            color: 'rgb(255, 0, 0)',
+            radius: 10,
+            shape: 'circle',
+          },
+        },
+      ]);
+    });
+    it('returns empty array if no legend', () => {
+      const legendInfo = factory
+        .create({
+          color: { value: 'black' },
+          shape: { value: 'square' },
+        })
+        .getLegendInformation([{ brand: 'Gucci' }, { brand: 'Prada' }]);
+
+      expect(stripFunction(legendInfo)).toEqual([]);
+    });
+  });
   describe('.hasLegend()', () => {
     it('returns true if has legend', () => {
       expect(encoder.hasLegend()).toBeTruthy();