From c158c80337e4b6a886a61b4d87510d311d75ae66 Mon Sep 17 00:00:00 2001
From: Cameron Yick <cameron.yick@datadoghq.com>
Date: Wed, 31 Jul 2024 14:21:46 -0400
Subject: [PATCH] fix(legend): use custom legend type throughout vega-lite
 compile functions

---
 src/compile/layer.ts            |  4 +++-
 src/compile/legend/assemble.ts  |  4 ++--
 src/compile/legend/component.ts | 20 +++++++++++++++++---
 src/compile/legend/encode.ts    |  6 +++---
 src/compile/model.ts            |  3 +--
 5 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/src/compile/layer.ts b/src/compile/layer.ts
index 933871c077..325014ff03 100644
--- a/src/compile/layer.ts
+++ b/src/compile/layer.ts
@@ -1,4 +1,4 @@
-import {Legend as VgLegend, NewSignal, SignalRef, Title as VgTitle} from 'vega';
+import {NewSignal, SignalRef, Title as VgTitle} from 'vega';
 import {array} from 'vega-util';
 import {Config} from '../config';
 import * as log from '../log';
@@ -13,6 +13,8 @@ import {parseLayerLayoutSize} from './layoutsize/parse';
 import {assembleLegends} from './legend/assemble';
 import {Model} from './model';
 import {assembleLayerSelectionMarks} from './selection/assemble';
+import {FullVgLegend as VgLegend} from './legend/component';
+
 import {UnitModel} from './unit';
 
 export class LayerModel extends Model {
diff --git a/src/compile/legend/assemble.ts b/src/compile/legend/assemble.ts
index 3732dac599..1137cff465 100644
--- a/src/compile/legend/assemble.ts
+++ b/src/compile/legend/assemble.ts
@@ -1,10 +1,10 @@
-import {Legend as VgLegend, LegendEncode} from 'vega';
+import {LegendEncode} from 'vega';
 import {Config} from '../../config';
 import {LEGEND_SCALE_CHANNELS} from '../../legend';
 import {keys, replaceAll, stringify, vals} from '../../util';
 import {isSignalRef, VgEncodeChannel, VgValueRef} from '../../vega.schema';
 import {Model} from '../model';
-import {LegendComponent} from './component';
+import {FullVgLegend as VgLegend, LegendComponent} from './component';
 import {mergeLegendComponent} from './parse';
 
 function setLegendEncode(
diff --git a/src/compile/legend/component.ts b/src/compile/legend/component.ts
index 849ed8660a..57ddec103d 100644
--- a/src/compile/legend/component.ts
+++ b/src/compile/legend/component.ts
@@ -1,10 +1,24 @@
-import {Legend as VgLegend} from 'vega';
+import {Legend as VgLegend, LegendType} from 'vega';
 import {NonPositionScaleChannel} from '../../channel';
 import {COMMON_LEGEND_PROPERTY_INDEX, LegendInternal} from '../../legend';
 import {Flag, keys} from '../../util';
 import {Split} from '../split';
 
-export type LegendComponentProps = VgLegend & {
+/** Update LegendType to support a third value (discrete), which is needed for discrete gradient legends
+ * https://github.com/vega/vega/blob/main/packages/vega-parser/src/parsers/legend.js */
+export type FullLegendType = LegendType | 'discrete';
+
+/** Update LegendType to support a third value (discrete), which is needed for discrete gradient legends
+ * https://github.com/vega/vega/blob/main/packages/vega-parser/src/parsers/legend.js */
+export type FullVgLegend = Omit<VgLegend, 'type'> & {
+  /**
+   * The type of legend to include. One of `"symbol"` for discrete symbol legends, `"gradient"` for a continuous color gradient, or `"discrete"` for a discrete color gradient.
+   * If gradient is used only the fill, stroke, and length scale parameters are considered. If unspecified, the type will be inferred based on the scale parameters used and their backing scale types.
+   * TODO: Remove FullLegendType override after base type in VgLegend is updated */
+  type?: FullLegendType;
+};
+
+export type LegendComponentProps = FullVgLegend & {
   labelExpr?: string;
   selections?: string[];
   disable?: boolean;
@@ -29,7 +43,7 @@ const LEGEND_COMPONENT_PROPERTY_INDEX: Flag<keyof LegendComponentProps> = {
 
 export const LEGEND_COMPONENT_PROPERTIES = keys(LEGEND_COMPONENT_PROPERTY_INDEX);
 
-export class LegendComponent extends Split<LegendComponentProps> {}
+export class LegendComponent extends Split<LegendComponentProps> { }
 
 export type LegendComponentIndex = Partial<Record<NonPositionScaleChannel, LegendComponent>>;
 
diff --git a/src/compile/legend/encode.ts b/src/compile/legend/encode.ts
index 334e0b0381..81887fadd9 100644
--- a/src/compile/legend/encode.ts
+++ b/src/compile/legend/encode.ts
@@ -1,4 +1,4 @@
-import {ColorValueRef, EncodeEntry, Gradient, LegendEncode, LegendType, SignalRef, SymbolEncodeEntry} from 'vega';
+import {ColorValueRef, EncodeEntry, Gradient, LegendEncode, SignalRef, SymbolEncodeEntry} from 'vega';
 import {array, isArray, stringValue} from 'vega-util';
 import {COLOR, NonPositionScaleChannel, OPACITY} from '../../channel';
 import {
@@ -19,14 +19,14 @@ import {formatCustomType, isCustomFormatType} from '../format';
 import * as mixins from '../mark/encode';
 import {STORE} from '../selection';
 import {UnitModel} from '../unit';
-import {LegendComponent} from './component';
+import {LegendComponent, FullLegendType} from './component';
 
 export interface LegendEncodeParams {
   fieldOrDatumDef: TypedFieldDef<string> | DatumDef;
   model: UnitModel;
   channel: NonPositionScaleChannel;
   legendCmpt: LegendComponent;
-  legendType: LegendType;
+  legendType: FullLegendType;
 }
 
 export const legendEncodeRules: {
diff --git a/src/compile/model.ts b/src/compile/model.ts
index 6dc99627d7..49347a2545 100644
--- a/src/compile/model.ts
+++ b/src/compile/model.ts
@@ -1,7 +1,6 @@
 import {
   AnchorValue,
   Axis as VgAxis,
-  Legend as VgLegend,
   NewSignal,
   Projection as VgProjection,
   Signal,
@@ -56,7 +55,7 @@ import {
   LayoutSizeType
 } from './layoutsize/component';
 import {assembleLegends} from './legend/assemble';
-import {LegendComponentIndex} from './legend/component';
+import {LegendComponentIndex, FullVgLegend as VgLegend} from './legend/component';
 import {parseLegend} from './legend/parse';
 import {assembleProjections} from './projection/assemble';
 import {ProjectionComponent} from './projection/component';