Add utilities for class-based LitTypes.

PiperOrigin-RevId: 459798788

lit_nlp/client/lib/lit_types.ts

@@ -15,8 +15,26 @@
  * limitations under the License.
  */
 
-type LitClass = 'LitType';
+// tslint:disable:no-new-decorators
+
+/**
+ * A dictionary of registered LitType names mapped to their constructor.
+ * LitTypes are added using the @registered decorator.
+ */
+export const REGISTRY: {[litType: string]: LitType} = {};
+// tslint:disable-next-line:no-any
+function registered(target: any) {
+  REGISTRY[target.name] = target;
+}
 
+const registryKeys = Object.keys(REGISTRY) as ReadonlyArray<string>;
+/**
+ * The types of all LitTypes in the registry, e.g.
+ * 'String' | 'TextSegment' ...
+ */
+export type LitName = typeof registryKeys[number];
+
+type LitClass = 'LitType';
 /**
  * Data classes used in configuring front-end components to describe
  * input data and model outputs.
@@ -43,13 +61,15 @@ export class LitType {
 /**
  * A string LitType.
  */
+@registered
 export class String extends LitType {
   override default: string = '';
 }
 
 /**
  * A scalar value, either a single float or int.
  */
+@registered
 export class Scalar extends LitType {
   override default: number = 0;
   min_val: number = 0;

lit_nlp/client/lib/lit_types_test.ts

@@ -1,13 +1,13 @@
 import 'jasmine';
+
 import * as litTypes from './lit_types';
 
 describe('lit types test', () => {
-
   it('creates a string', () => {
     const testString = new litTypes.String();
-    testString.default = "string value";
+    testString.default = 'string value';
 
-    expect(testString.default).toBe("string value");
+    expect(testString.default).toBe('string value');
     expect(testString.required).toBe(true);
 
     expect(testString instanceof litTypes.String).toBe(true);

lit_nlp/client/lib/lit_types_utils.ts

@@ -0,0 +1,86 @@
+import {Spec} from '../lib/types';
+
+import {LitName, LitType, REGISTRY} from './lit_types';
+
+
+/**
+ * Creates and returns a new LitType instance.
+ * @param typeName: The name of the desired LitType.
+ * @param constructorParams: A dictionary of properties to set on the LitType.
+ * For example, {'show_in_data_table': true}.
+ */
+export function createLitType(
+    typeName: LitName,
+    constructorParams: {[key: string]: unknown} = {}){
+  const litType = REGISTRY[typeName];
+
+  // tslint:disable-next-line:no-any
+  const newType = new (litType as any)();
+  newType.__name__ = typeName;
+  newType.__mro__ = getMethodResolutionOrder(newType);
+
+  for (const key in constructorParams) {
+    if (key in newType) {
+      newType[key] = constructorParams[key];
+    } else {
+      throw new Error(
+          `Attempted to set unrecognized property ${key} on ${newType}.`);
+    }
+  }
+
+  return newType;
+}
+
+/**
+ * Returns the method resolution order for a given litType.
+ * This is for compatability with references to non-class-based LitTypes,
+ * and should match the Python class hierarchy.
+ */
+export function getMethodResolutionOrder(litType: LitType): string[] {
+  const mro: string[] = [];
+
+  // TODO(b/162269499): Remove this method after we replace the old LitType.
+  let object = Object.getPrototypeOf(litType);
+  while (object) {
+    mro.push(object.constructor.name);
+    object = Object.getPrototypeOf(object);
+  }
+
+  return mro;
+}
+
+/**
+ * Returns whether the litType is a subtype of any of the typesToFind.
+ * @param litType: The LitType to check.
+ * @param typesToFind: Either a single or list of parent LitType candidates.
+ */
+export function isLitSubtype(
+    litType: LitType, typesToFind: LitName|LitName[]) {
+  if (litType == null) return false;
+
+  if (typeof typesToFind === 'string') {
+    typesToFind = [typesToFind];
+  }
+
+  for (const typeName of typesToFind) {
+    // tslint:disable-next-line:no-any
+    const registryType : any = REGISTRY[typeName];
+
+    if (litType instanceof registryType) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/**
+ * Returns all keys in the given spec that are subtypes of the typesToFind.
+ * @param spec: A Spec object.
+ * @param typesToFind: Either a single or list of parent LitType candidates.
+ */
+export function findSpecKeys(
+    spec: Spec, typesToFind: LitName|LitName[]): string[] {
+  return Object.keys(spec).filter(
+      key => isLitSubtype(
+          spec[key] as LitType, typesToFind));
+}

lit_nlp/client/lib/lit_types_utils_test.ts

@@ -0,0 +1,106 @@
+import 'jasmine';
+
+import {Spec} from '../lib/types';
+
+import * as litTypes from './lit_types';
+import * as litTypesUtils from './lit_types_utils';
+
+
+describe('createLitType test', () => {
+  it('creates a type as expected', () => {
+    const expected = new litTypes.Scalar();
+    expected.__name__ = 'Scalar';
+    expected.__mro__ = ['Scalar', 'LitType', 'Object'];
+    expected.show_in_data_table = false;
+
+    const result = litTypesUtils.createLitType('Scalar');
+    expect(result).toEqual(expected);
+  });
+
+  it('creates with constructor params', () => {
+    const expected = new litTypes.String();
+    expected.__name__ = 'String';
+    expected.__mro__ = ['String', 'LitType', 'Object'];
+    expected.default = 'foo';
+    expected.show_in_data_table = true;
+
+    const result = litTypesUtils.createLitType(
+        'String', {'show_in_data_table': true, 'default': 'foo'});
+    expect(result).toEqual(expected);
+  });
+
+  it('creates a modifiable lit type', () => {
+    const result = litTypesUtils.createLitType('Scalar');
+    result.min_val = 5;
+    expect(result.min_val).toEqual(5);
+  });
+
+  it('handles invalid constructor params', () => {
+    expect(() => litTypesUtils.createLitType('String', {
+      'notAStringParam': true
+    })).toThrowError();
+  });
+
+  it('populates mro', () => {
+    let testType = new litTypes.String();
+    expect(litTypesUtils.getMethodResolutionOrder(testType)).toEqual([
+      'String', 'LitType', 'Object'
+    ]);
+
+    testType = new litTypes.LitType();
+    expect(litTypesUtils.getMethodResolutionOrder(testType)).toEqual([
+      'LitType', 'Object'
+    ]);
+  });
+
+  it('handles invalid names', () => {
+    expect(() => litTypesUtils.createLitType('notLitType')).toThrowError();
+  });
+});
+
+describe('isLitSubtype test', () => {
+  it('checks is lit subtype', () => {
+    const testType = new litTypes.String();
+    expect(litTypesUtils.isLitSubtype(testType, 'String')).toBe(true);
+    expect(litTypesUtils.isLitSubtype(testType, ['String'])).toBe(true);
+    expect(litTypesUtils.isLitSubtype(testType, ['Scalar'])).toBe(false);
+
+    // LitType is not a subtype of LitType.
+    expect(() => litTypesUtils.isLitSubtype(testType, 'LitType'))
+        .toThrowError();
+    expect(() => litTypesUtils.isLitSubtype(testType, ['NotAType']))
+        .toThrowError();
+  });
+});
+
+
+describe('findSpecKeys test', () => {
+  // TODO(cjqian): Add original utils_test test after adding more types.
+  const spec: Spec = {
+    'scalar_foo': new litTypes.Scalar(),
+    'segment': new litTypes.String(),
+    'generated_text': new litTypes.String(),
+  };
+
+
+  it('finds all spec keys that match the specified types', () => {
+    // Key is in spec.
+    expect(litTypesUtils.findSpecKeys(spec, 'String')).toEqual([
+      'segment', 'generated_text'
+    ]);
+
+    // Keys are in spec.
+    expect(litTypesUtils.findSpecKeys(spec, ['String', 'Scalar'])).toEqual([
+      'scalar_foo', 'segment', 'generated_text'
+    ]);
+  });
+
+  it('handles empty spec keys', () => {
+    expect(litTypesUtils.findSpecKeys(spec, [])).toEqual([]);
+  });
+
+  it('handles invalid spec keys', () => {
+    expect(() => litTypesUtils.findSpecKeys(spec, '')).toThrowError();
+    expect(() => litTypesUtils.findSpecKeys(spec, 'NotAType')).toThrowError();
+  });
+});

lit_nlp/client/services/data_service.ts

@@ -19,6 +19,7 @@
 import {action, computed, observable, reaction} from 'mobx';
 
 import {BINARY_NEG_POS, ColorRange} from '../lib/colors';
+import {createLitType} from '../lib/lit_types_utils';
 import {ClassificationResults, IndexedInput, LitName, LitType, RegressionResults} from '../lib/types';
 import {findSpecKeys, isLitSubtype, mapsContainSame} from '../lib/utils';
 
@@ -213,7 +214,7 @@ export class DataService extends LitService {
           (result: RegressionResults) => result[key].error);
       const sqErrors = regressionResults.map(
           (result: RegressionResults) => result[key].squared_error);
-      const dataType = this.appState.createLitType('Scalar', false);
+      const dataType = createLitType('Scalar');
       const source = `${REGRESSION_SOURCE_PREFIX}:${model}`;
       this.addColumnFromList(scores, data, scoreFeatName, dataType, source);
       if (output[key].parent != null) {

lit_nlp/client/services/data_service_test.ts

@@ -20,6 +20,7 @@
 import 'jasmine';
 
 import {LitApp} from '../core/app';
+import {createLitType} from '../lib/lit_types_utils';
 import {mockMetadata} from '../lib/testing_utils';
 import {IndexedInput} from '../lib/types';
 import {ApiService, AppState, ClassificationService, SettingsService, StatusService} from '../services/services';
@@ -128,7 +129,7 @@ describe('DataService test', () => {
   it('has correct columns', () => {
     expect(dataService.cols.length).toBe(0);
 
-    const dataType = appState.createLitType('Scalar');
+    const dataType = createLitType('Scalar');
     const getValueFn = () => 1;
     const dataMap: ColumnData = new Map();
     for (let i = 0; i < appState.currentInputData.length; i++) {