Merge isLitSubtype and findSpecKeys implementations from `lit_typ…

…e_utils` and `utils`, now that we are using class-based LitTypes throughout the codebase.

PiperOrigin-RevId: 463642302

lit_nlp/client/lib/lit_types.ts

@@ -44,6 +44,7 @@ type ScoredTextCandidates = Array<[text: string, score: number|null]>;
  * Data classes used in configuring front-end components to describe
  * input data and model outputs.
  */
+@registered
 export class LitType {
   // tslint:disable:enforce-name-casing
   __class__: LitClass|'type' = 'LitType';

lit_nlp/client/lib/lit_types_utils.ts

@@ -18,11 +18,8 @@
 // For consistency with types.ts.
 // tslint:disable: enforce-name-casing
 
-import {Spec} from '../lib/types';
-
-import {LitName, LitType, REGISTRY} from './lit_types';
-import {LitMetadata} from './types';
-
+import {LitName, REGISTRY} from './lit_types';
+import {LitMetadata, Spec} from './types';
 
 /**
  * Creates and returns a new LitType instance.
@@ -36,7 +33,6 @@ export function createLitType(
   // tslint:disable-next-line:no-any
   const newType = new (litType as any)();
   newType.__name__ = typeName;
-  newType.__mro__ = getMethodResolutionOrder(newType);
 
   // Excluded properties are passed through in the Python serialization
   // of LitTypes and can be ignored by the frontend.
@@ -107,56 +103,3 @@ export function deserializeLitTypesInLitMetadata(metadata: LitMetadata):
   metadata.littypes = deserializeLitTypesInSpec(metadata.littypes);
   return metadata;
 }
-
-
-/**
- * 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], typesToFind));
-}

lit_nlp/client/lib/lit_types_utils_test.ts

@@ -1,7 +1,5 @@
 import 'jasmine';
 
-import {Spec} from '../lib/types';
-
 import * as litTypes from './lit_types';
 import * as litTypesUtils from './lit_types_utils';
 
@@ -10,7 +8,6 @@ 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');
@@ -21,7 +18,6 @@ describe('createLitType test', () => {
   it('creates with constructor params', () => {
     const expected = new litTypes.StringLitType();
     expected.__name__ = 'StringLitType';
-    expected.__mro__ = ['StringLitType', 'LitType', 'Object'];
     expected.default = 'foo';
     expected.show_in_data_table = true;
 
@@ -56,39 +52,11 @@ describe('createLitType test', () => {
     expect(categoryLabel.vocab).toEqual(vocab);
   });
 
-  it('populates mro', () => {
-    let testType = new litTypes.StringLitType();
-    expect(litTypesUtils.getMethodResolutionOrder(testType)).toEqual([
-      'StringLitType', '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.StringLitType();
-    expect(litTypesUtils.isLitSubtype(testType, 'StringLitType')).toBe(true);
-    expect(litTypesUtils.isLitSubtype(testType, ['StringLitType'])).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('deserializeLitTypesInSpec test', () => {
   // TODO(b/162269499): Add test for deserializeLitTypesInLitMetadata.
   const testSpec = {
@@ -121,36 +89,3 @@ describe('deserializeLitTypesInSpec test', () => {
         .toBe(true);
   });
 });
-
-
-describe('findSpecKeys test', () => {
-  // TODO(cjqian): Add original litTypesUtils.test test after adding more types.
-  const spec: Spec = {
-    'scalar_foo': new litTypes.Scalar(),
-    'segment': new litTypes.StringLitType(),
-    'generated_text': new litTypes.StringLitType(),
-  };
-
-
-  it('finds all spec keys that match the specified types', () => {
-    // Key is in spec.
-    expect(litTypesUtils.findSpecKeys(spec, 'StringLitType')).toEqual([
-      'segment', 'generated_text'
-    ]);
-
-    // Keys are in spec.
-    expect(litTypesUtils.findSpecKeys(spec, [
-      'StringLitType', '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/lib/utils.ts

@@ -25,7 +25,7 @@ import {html, TemplateResult} from 'lit';
 import {unsafeHTML} from 'lit/directives/unsafe-html.js';
 
 import {marked} from 'marked';
-import {LitName, LitType, LitTypeWithParent, MulticlassPreds} from './lit_types';
+import {LitName, LitType, LitTypeWithParent, MulticlassPreds, REGISTRY} from './lit_types';
 import {FacetMap, ModelInfoMap, Spec} from './types';
 
 /** Calculates the mean for a list of numbers */
@@ -91,40 +91,42 @@ export function mapsContainSame<T>(mapA: Map<string, T>, mapB: Map<string, T>) {
 }
 
 /**
- * Check if a spec field (LitType) is an instance of one or more type names.
- * This is analogous to using isinstance(litType, typesToFind) in Python,
- * and relies on exporting the Python class hierarchy in the __mro__ field.
+ * 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[]) {
-  // TODO(lit-dev): figure out why this is occasionally called on an invalid
-  // spec. Likely due to skew between keys and specs in specific modules when
-  // dataset is changed, but worth diagnosing to make sure this doesn't mask a
-  // larger issue.
+export function isLitSubtype(
+    litType: LitType, typesToFind: LitName|LitName[]) {
   if (litType == null) return false;
 
   if (typeof typesToFind === 'string') {
     typesToFind = [typesToFind];
   }
+
   for (const typeName of typesToFind) {
-    if (litType.__mro__.includes(typeName)) {
+    // tslint:disable-next-line:no-any
+    const registryType : any = REGISTRY[typeName];
+
+    if (litType instanceof registryType) {
       return true;
     }
   }
   return false;
 }
 
+
 /**
- * Find all keys from the spec which match any of the specified types.
+ * 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[] {
-  if (typeof typesToFind === 'string') {
-    typesToFind = [typesToFind];
-  }
   return Object.keys(spec).filter(
-      key => isLitSubtype(spec[key], typesToFind as LitName[]));
+      key => isLitSubtype(spec[key], typesToFind));
 }
 
+
 /**
  * Return a new object with the selected keys from the old one.
  */

lit_nlp/client/lib/utils_test.ts

@@ -141,6 +141,16 @@ describe('arrayContainsSame test', () => {
 });
 
 describe('isLitSubtype test', () => {
+  it('checks is lit subtype', () => {
+    const testType = createLitType('StringLitType');
+    expect(utils.isLitSubtype(testType, 'StringLitType')).toBe(true);
+    expect(utils.isLitSubtype(testType, ['StringLitType'])).toBe(true);
+    expect(utils.isLitSubtype(testType, ['Scalar'])).toBe(false);
+
+    expect(() => utils.isLitSubtype(testType, ['NotAType']))
+        .toThrowError();
+  });
+
   it('finds a subclass', () => {
     const spec: Spec = {
       'score': createLitType('RegressionScore'),
@@ -193,6 +203,16 @@ describe('findSpecKeys test', () => {
       'score', 'score2', 'scalar_foo'
     ]);
   });
+
+  it('handles empty spec keys', () => {
+    expect(utils.findSpecKeys(spec, [])).toEqual([]);
+  });
+
+  it('handles invalid spec keys', () => {
+    expect(() => utils.findSpecKeys(spec, '')).toThrowError();
+    expect(() => utils.findSpecKeys(spec, 'NotAType')).toThrowError();
+  });
+
 });
 
 
@@ -564,4 +584,4 @@ describe('linearSpace test', () => {
     result = [];
     expect(utils.linearSpace(minValue, maxValue, numSteps)).toEqual(result);
   });
-});
+});