feat: Expose expression functions (#935)

Co-authored-by: Kanit Wongsuphasawat <[email protected]>

README.md

@@ -159,6 +159,7 @@ var opt = {
 
   formatLocale: ...,
   timeFormatLocale: ...,
+  expressionFunctions: ...,
 
   ast: ...,
   expr: ...,
@@ -198,6 +199,7 @@ var opt = {
 | `downloadFileName` | String   | Sets the file name (default: `visualization`) for charts downloaded using the `png` or `svg` action. |
 | `formatLocale` | Object   | Sets the default locale definition for number formatting. See the [d3-format locale collection](https://github.com/d3/d3-format/tree/master/locale) for definition files for a variety of languages. Note that this is a global setting. |
 | `timeFormatLocale` | Object   | Sets the default locale definition for date/time formatting. See the [d3-time-format locale collection](https://github.com/d3/d3-time-format/tree/master/locale) for definition files for a variety of languages. Note that this is a global setting. |
+| `expressionFunctions` | Object   | Sets custom expression functions. Maps a function name to a JavaScript `function`, or an Object with the `fn`, and `visitor` parameters. See [Vega Expression Functions](https://vega.github.io/vega/docs/api/extensibility/#expressionFunction) for more information. |
 | `ast` | Boolean   | Generate an [Abstract Syntax Tree (AST)](https://en.wikipedia.org/wiki/Abstract_syntax_tree) instead of expressions and use an interpreter instead of native evaluation. While the interpreter is slower, it adds support for Vega expressions that are [Content Security Policy (CSP)](https://developer.mozilla.org/en-US/docs/Web/HTTP/CSP)-compliant. |
 | `expr` | Object   | Custom Vega Expression interpreter. |
 | `viewClass` | Class         | Class which extends [Vega `View`](https://vega.github.io/vega/docs/api/view/#view) for custom rendering. |

src/embed.ts

@@ -25,7 +25,7 @@ import * as themes from 'vega-themes';
 import {Handler, Options as TooltipOptions} from 'vega-tooltip';
 import post from './post';
 import embedStyle from './style';
-import {Config, Mode} from './types';
+import {Config, ExpressionFunction, Mode} from './types';
 import {mergeDeep} from './util';
 import pkg from '../package.json';
 
@@ -91,6 +91,7 @@ export interface EmbedOptions<S = string, R = Renderers> {
   downloadFileName?: string;
   formatLocale?: Record<string, unknown>;
   timeFormatLocale?: Record<string, unknown>;
+  expressionFunctions?: ExpressionFunction;
   ast?: boolean;
   expr?: typeof expressionInterpreter;
   viewClass?: typeof View;
@@ -163,8 +164,7 @@ export function guessMode(spec: VisualizationSpec, providedMode?: Mode): Mode {
     const parsed = schemaParser(spec.$schema);
     if (providedMode && providedMode !== parsed.library) {
       console.warn(
-        `The given visualization spec is written in ${NAMES[parsed.library]}, but mode argument sets ${
-          NAMES[providedMode] ?? providedMode
+        `The given visualization spec is written in ${NAMES[parsed.library]}, but mode argument sets ${NAMES[providedMode] ?? providedMode
         }.`
       );
     }
@@ -352,6 +352,18 @@ async function _embed(
     vega.timeFormatLocale(opts.timeFormatLocale);
   }
 
+  // Set custom expression functions
+  if (opts.expressionFunctions) {
+    for (const name in opts.expressionFunctions) {
+      const expressionFunction = opts.expressionFunctions[name];
+      if ('fn' in expressionFunction) {
+        vega.expressionFunction(name, expressionFunction.fn, expressionFunction['visitor']);
+      } else if (expressionFunction instanceof Function) {
+        vega.expressionFunction(name, expressionFunction);
+      }
+    }
+  }
+
   const {ast} = opts;
 
   // Do not apply the config to Vega when we have already applied it to Vega-Lite.
@@ -384,7 +396,7 @@ async function _embed(
     const handler = isTooltipHandler(opts.tooltip)
       ? opts.tooltip
       : // user provided boolean true or tooltip options
-        new Handler(opts.tooltip === true ? {} : opts.tooltip).call;
+      new Handler(opts.tooltip === true ? {} : opts.tooltip).call;
 
     view.tooltip(handler);
   }

src/types.ts

@@ -4,6 +4,8 @@ import {Config as VlConfig} from 'vega-lite';
 export type Mode = 'vega' | 'vega-lite';
 export type Config = VlConfig | VgConfig;
 
+export type ExpressionFunction = Record<string, any | {fn: any; visitor?: any}>;
+
 export interface MessageData {
   spec: string;
   file?: unknown;

test/embed.test.ts

@@ -13,6 +13,23 @@ const vlSpec: TopLevelSpec = {
 
 const vgSpec = compile(vlSpec).spec;
 
+const vlSpecCustomFunction: TopLevelSpec = {
+  data: {values: [1, 2, 3]},
+  encoding: {
+    y: {
+      axis: {
+        format: '',
+        formatType: 'simpleFunction',
+      },
+    },
+  },
+  mark: 'point',
+  transform: [
+    {calculate: 'simpleFunction()', as: 'result1'},
+    {calculate: 'functionWithVisitor()', as: 'result2'},
+  ],
+};
+
 test('embed returns result', async () => {
   const el = document.createElement('div');
   const result = await embed(el, vlSpec);
@@ -242,6 +259,43 @@ test('can set locale', async () => {
   expect(result).toBeTruthy();
 });
 
+test('throws error when expressionFunction does not exist', async () => {
+  const el = document.createElement('div');
+
+  const getErrorFromEmbed = async () => {
+    try {
+      await embed(el, vlSpecCustomFunction);
+
+      throw Error('No Thrown Error');
+    } catch (e: any) {
+      return e;
+    }
+  };
+
+  const error = await getErrorFromEmbed();
+  expect(error.message).toBe('Unrecognized function: simpleFunction');
+});
+
+test('can set and use expressionFunctions', async () => {
+  const el = document.createElement('div');
+  const result = await embed(el, vlSpecCustomFunction, {
+    expressionFunctions: {
+      simpleFunction: () => {
+        return 'test';
+      },
+      functionWithVisitor: {
+        fn: () => {
+          return 'test';
+        },
+        visitor: () => {
+          return 'test';
+        },
+      },
+    },
+  });
+  expect(result).toBeTruthy();
+});
+
 test('can set tooltip theme', async () => {
   const el = document.createElement('div');
   const result = await embed(el, vlSpec, {