Improve Typing for Aesthetics and Transforms

src/Aesthetic.ts

@@ -7,6 +7,8 @@ import {
   scaleOrdinal,
   scaleSequential, scaleSequentialLog, 
   scaleSequentialPow,
+  ScaleContinuousNumeric,
+  ScaleIdentity,
 } from 'd3-scale';
 import { rgb } from 'd3-color';
 import * as d3Chromatic from 'd3-scale-chromatic';
@@ -16,7 +18,8 @@ import type {
   TextureSet
 } from './AestheticSet';
 import { isOpChannel, isLambdaChannel,
-  isConstantChannel } from './types';
+  isConstantChannel, 
+  Transform} from './types';
 import type {
   OpChannel, LambdaChannel,
   Channel, BasicChannel, ConstantChannel,
@@ -27,7 +30,10 @@ import { Vector } from 'apache-arrow';
 import { StructRowProxy } from 'apache-arrow/row/struct';
 import { QuadTile } from './tile';
 
-const scales : { [key : string] : Function } = {
+const scales: Record<
+  Transform,
+  (range?: Iterable<number>) => ScaleContinuousNumeric<number, number, never> | ScaleIdentity<never>
+> = {
   sqrt: scaleSqrt,
   log: scaleLog,
   linear: scaleLinear,
@@ -98,19 +104,17 @@ function okabe() {
 
 okabe();
 
-type Transform = 'log' | 'sqrt' | 'linear' | 'literal';
-
 abstract class Aesthetic {
   public abstract default_range : [number, number];
   public abstract default_constant : number | [number, number, number];
-  public abstract default_transform : 'log' | 'sqrt' | 'linear' | 'literal';
+  public abstract default_transform : Transform
   public scatterplot : Scatterplot;
   public field: string | null = null;
   public regl : Regl;
   public _texture_buffer : Float32Array | Uint8Array | null = null;
   public _domain : [number, number];
   public _range : [number, number] | Uint8Array;
-  public _transform : 'log' | 'sqrt' | 'linear' | 'literal' | undefined;
+  public _transform : Transform | undefined;
   public dataset : QuadtileSet;
   public partner : Aesthetic | null = null;
   public _textures : Record<string, Texture2D> = {};
@@ -149,14 +153,11 @@ abstract class Aesthetic {
   }
 
   get transform() {    
-    if (this._transform) return this._transform;
-    return this.default_transform;
+    return this._transform || this.default_transform;
   }
 
   set transform(transform) {
-
     this._transform = transform;
-
   }
 
   get scale() {
@@ -165,9 +166,9 @@ abstract class Aesthetic {
       return r.charAt(0).toUpperCase() + r.slice(1);
     }
 
-    let scale = scales[this.transform]()
-      .domain(this.domain)
-      .range(this.range);
+    const domain = scales[this.transform]()
+      .domain(this.domain) as ScaleContinuousNumeric<number, number, never>;
+    let scale = domain.range(this.range);
 
     const range = this.range;
 
@@ -688,15 +689,16 @@ export const dimensions = {
   y : Y
 };
 
-type concrete_aesthetics = X | 
-Y |
-Size | 
-Jitter_speed |
-Jitter_radius | 
-Color | 
-Filter;
+type ConcreteAesthetic = 
+  | X
+  | Y
+  | Size
+  | Jitter_speed
+  | Jitter_radius
+  | Color
+  | Filter;
 
-export abstract class StatefulAesthetic<T extends concrete_aesthetics> {
+export abstract class StatefulAesthetic<T extends ConcreteAesthetic> {
   // An aesthetic that tracks two states--current and last.
   // The point is to handle transitions.
   // It might make sense to handle more than two states, but there are
@@ -775,7 +777,7 @@ class StatefulColor extends StatefulAesthetic<Color> { get Factory() { return Co
 class StatefulFilter extends StatefulAesthetic<Filter> { get Factory() { return Filter; }} 
 class StatefulFilter2 extends StatefulAesthetic<Filter> { get Factory() { return Filter; }}
 
-export const stateful_aesthetics : Record<string, StatefulAesthetic<typeof Aesthetic>> = {
+export const stateful_aesthetics = {
   x : StatefulX,
   x0 : StatefulX0,
   y : StatefulY,
@@ -787,6 +789,7 @@ export const stateful_aesthetics : Record<string, StatefulAesthetic<typeof Aesth
   filter : StatefulFilter,
   filter2 : StatefulFilter2
 };
+export type StatefulAestheticSet = typeof stateful_aesthetics;
 
 function parseLambdaString(lambdastring : string) {
   // Materialize an arrow function from its string.
@@ -856,4 +859,4 @@ function lambda_to_function(input : LambdaChannel) : (d : any) => number {
   //@ts-ignore
   const func : (d : any) => number = new Function(arg, code);
   return func;
-}
+}

src/AestheticSet.ts

@@ -1,20 +1,25 @@
 /* eslint-disable no-param-reassign */
 import type { Regl, Texture2D } from 'regl';
 import { 
+  StatefulAestheticSet,
   stateful_aesthetics,
 } from './Aesthetic';
 import type Scatterplot from './deepscatter';
 import type QuadtreeRoot from './tile';
 import type { Encoding } from './types';
 import type { StatefulAesthetic } from './Aesthetic';
 
+type StatefulAestheticStore = {
+  [K in keyof StatefulAestheticSet]?: InstanceType<StatefulAestheticSet[K]>
+};
+
 export class AestheticSet {
   public tileSet : QuadtreeRoot;
   public scatterplot : Scatterplot;
   public regl : Regl;
   public encoding : Encoding;
   public position_interpolation : boolean;
-  private store : Record<string, StatefulAesthetic<any>>;
+  private store : StatefulAestheticStore;
   public aesthetic_map : TextureSet;
   constructor(scatterplot : Scatterplot, regl : Regl, tileSet : QuadtreeRoot) {
     this.scatterplot = scatterplot;
@@ -26,17 +31,21 @@ export class AestheticSet {
     return this;
   }
 
-  public dim(aesthetic : string) {
+  public dim<K extends keyof StatefulAestheticSet>(aesthetic : K): NonNullable<StatefulAestheticStore[K]> {
     // Returns the stateful aesthetic corresponding to the given aesthetic.
     if (this.store[aesthetic]) {
-      return this.store[aesthetic];
+      const storedAesthetic = this.store[aesthetic]!;
+      return storedAesthetic;
     }
-    if (stateful_aesthetics[aesthetic] !== undefined) {
-      this.store[aesthetic] = new stateful_aesthetics[aesthetic](
+
+    // if (stateful_aesthetics[aesthetic] !== undefined) {
+    if (aesthetic in stateful_aesthetics) {
+      const newAesthetic = (new stateful_aesthetics[aesthetic](
         this.scatterplot, this.regl, this.tileSet, 
         this.aesthetic_map
-      );
-      return this.store[aesthetic];
+      )) as  NonNullable<StatefulAestheticStore[K]>;
+      this.store[aesthetic] = newAesthetic;
+      return newAesthetic;
     }
     throw new Error(`Unknown aesthetic ${aesthetic}`);
   }
@@ -204,4 +213,4 @@ export class TextureSet {
     });
     return this._color_texture;
   }
-}
+}

src/Dataset.ts

@@ -45,7 +45,7 @@ export abstract class Dataset<T extends Tile> {
 
   map<U>(callback : (tile: T) => U, after = false) : U[] {
     const results : U[] = [];
-    this.visit((d : T) => { results.push(callback(d)); }, after = after);
+    this.visit((d : T) => { results.push(callback(d)); }, after);
     return results;
   }
 
@@ -248,4 +248,4 @@ function check_overlap(tile : Tile, bbox : Rectangle) : number {
     return disqualify;
   }
   return area(intersection) / area(bbox);
-}
+}

src/regl_rendering.ts

@@ -126,7 +126,7 @@ export class ReglRenderer extends Renderer {
     const { prefs } = this;
     const { transform } = this.zoom;
     const { aes_to_buffer_num, buffer_num_to_variable, variable_to_buffer_num } = this.allocate_aesthetic_buffers();
-    console.log(prefs.arrow_table);
+    // console.log(prefs.arrow_table);
     const props = {
     // Copy the aesthetic as a string.
       aes: { encoding: this.aes.encoding },

src/types.ts

@@ -36,6 +36,7 @@ export type ConstantChannel = {
   constant : number;
 };
 
+export type Transform = 'log' | 'linear' | 'sqrt' | 'literal';
 
 /**
  * A channel represents the information necessary to map a single dimension
@@ -56,7 +57,7 @@ export interface BasicChannel {
    * 'literal' maps in the implied dataspace set by 'x', 'y', while
    * 'linear' transforms the data by the range and domain.
   */
-  transform? : 'log' | 'linear' | 'sqrt' | 'literal';
+  transform? : Transform;
   // The domain over which the data extends 
   domain? : [number, number];
   // The range into which to map the data.
@@ -167,4 +168,4 @@ export function isLambdaChannel(input: Channel): input is LambdaChannel {
 
 export function isConstantChannel(input: Channel): input is ConstantChannel {
   return (input as ConstantChannel).constant !== undefined;
-}
+}