Selected Color by value synced to URL param

PiperOrigin-RevId: 480715945

lit_nlp/client/core/app.ts

@@ -55,6 +55,7 @@ export class LitApp {
     const [selectionService, pinnedSelectionService] =
         this.getServiceArray(SelectionService);
     const urlService = this.getService(UrlService);
+    const colorService = this.getService(ColorService);
 
     // Load the app metadata before any further initialization
     appState.metadata = await apiService.getInfo();
@@ -67,8 +68,8 @@ export class LitApp {
     }
 
     // Sync app state based on URL search params
-    urlService.syncStateToUrl(
-        appState, modulesService, selectionService, pinnedSelectionService);
+    urlService.syncStateToUrl(appState, modulesService, selectionService,
+        pinnedSelectionService, colorService);
 
     // Initialize the rest of the app state
     await appState.initialize();
@@ -156,8 +157,7 @@ export class LitApp {
     const dataService = new DataService(
         appState, classificationService, apiService, settingsService);
     const groupService = new GroupService(appState, dataService);
-    const colorService = new ColorService(
-        appState, groupService, dataService);
+    const colorService = new ColorService(groupService, dataService);
     const focusService = new FocusService(selectionService);
 
     // Populate the internal services map for dependency injection

lit_nlp/client/core/main_toolbar.ts

@@ -161,7 +161,7 @@ export class LitMainMenu extends MobxLitElement {
         displayIcon: isSelected,
         menu: [],
         onClick: () => {
-          this.colorService.selectedColorOption = option;
+          this.colorService.selectedColorOptionName = option.name;
         },
         disabled: false,
       };

lit_nlp/client/services/color_service.ts

@@ -17,15 +17,15 @@
 
 // tslint:disable:no-new-decorators
 import * as d3 from 'd3';
-import {computed, observable, reaction} from 'mobx';
+import {computed, observable} from 'mobx';
 
 import {CATEGORICAL_NORMAL, CONTINUOUS_SIGNED_LAB, CONTINUOUS_UNSIGNED_LAB, DEFAULT, MULTIHUE_CONTINUOUS} from '../lib/colors';
 import {ColorOption, D3Scale, IndexedInput} from '../lib/types';
 
 import {DataService} from './data_service';
 import {GroupService} from './group_service';
 import {LitService} from './lit_service';
-import {AppState} from './state_service';
+import {ColorObservedByUrlService, UrlConfiguration} from './url_service';
 
 /** Color map for salience maps. */
 export abstract class SalienceCmap {
@@ -97,15 +97,12 @@ export class SignedSalienceCmap extends SalienceCmap {
 /**
  * A singleton class that handles all coloring options.
  */
-export class ColorService extends LitService {
+export class ColorService extends LitService implements
+    ColorObservedByUrlService {
   constructor(
-      private readonly appState: AppState,
       private readonly groupService: GroupService,
       private readonly dataService: DataService) {
     super();
-    reaction(() => this.appState.currentModels, currentModels => {
-      this.reset();
-    });
   }
 
   private readonly defaultColor = DEFAULT;
@@ -118,7 +115,11 @@ export class ColorService extends LitService {
 
   // Name of selected feature to color datapoints by, or default not coloring by
   // features.
-  @observable selectedColorOption = this.defaultOption;
+  @observable mySelectedColorOption = this.defaultOption;
+  // It's used for the url service. When urlService.syncStateToUrl is invoked,
+  // colorableOptions are not available. There, this variable is used to
+  // preserve the url param value entered by users.
+  @observable selectedColorOptionName: string = '';
 
   // All variables that affect color settings, so clients can listen for when
   // they may need to rerender.
@@ -129,6 +130,19 @@ export class ColorService extends LitService {
     ];
   }
 
+  // Return the selectedColorOption based on the selectedColorOptionName
+  @computed
+  get selectedColorOption() {
+    if (this.colorableOptions.length === 0 ||
+        this.selectedColorOptionName.length === 0) {
+      return this.defaultOption;
+    } else {
+      const options = this.colorableOptions.filter(
+          option => option.name === this.selectedColorOptionName);
+      return options.length ? options[0] : this.defaultOption;
+    }
+  }
+
   @computed
   get colorableOptions() {
     // TODO(b/156100081): Get proper reactions on data service columns.
@@ -192,10 +206,8 @@ export class ColorService extends LitService {
     return this.selectedColorOption.scale(val) || this.defaultColor;
   }
 
-  /**
-   * Reset stored info. Used when active models change.
-   */
-  reset() {
-    this.selectedColorOption = this.defaultOption;
+  // Set color option based on the URL configuration
+  setUrlConfiguration(urlConfiguration: UrlConfiguration) {
+    this.selectedColorOptionName = urlConfiguration.colorBy ?? '';
   }
 }

lit_nlp/client/services/color_service_test.ts

@@ -56,8 +56,7 @@ describe('Color service test', () => {
 
   const app = new LitApp();
   const colorService = new ColorService(
-      app.getService(AppState), mockGroupService,
-      app.getService(DataService));
+      mockGroupService, app.getService(DataService));
 
   it('Tests colorableOption', () => {
     const opts = colorService.colorableOptions;
@@ -106,19 +105,16 @@ describe('Color service test', () => {
 
     // When the settings are updated, getDatapointColor() should reflect the
     // update.
-    colorService.selectedColorOption = colorService.colorableOptions[0];
+    colorService.selectedColorOptionName =
+        colorService.colorableOptions[0].name;
     color = colorService.getDatapointColor(mockInput);
     expect(color).toEqual(CATEGORICAL_NORMAL[1]);
 
     // Updating to a numerical color scheme.
-    colorService.selectedColorOption = colorService.colorableOptions[2];
+    colorService.selectedColorOptionName =
+        colorService.colorableOptions[2].name;
     color = colorService.getDatapointColor(mockInput);
     expect(color).toEqual('rgb(51, 138, 163)');
-
-    // After resetting, getDatapointColor() should reset.
-    colorService.reset();
-    color = colorService.getDatapointColor(mockInput);
-    expect(color).toEqual(DEFAULT);
   });
 
   it('provides color map classes for salience viz', () => {

lit_nlp/client/services/url_service.ts

@@ -46,6 +46,7 @@ export class UrlConfiguration {
   /** Path to load a new dataset from, on pageload. */
   newDatasetPath?: string;
   documentationOpen?: boolean;
+  colorBy?: string;
 }
 
 /**
@@ -88,6 +89,14 @@ export interface SelectionObservedByUrlService {
   selectIds: (ids: string[], user: ServiceUser) => void;
 }
 
+/**
+ * Interface describing how the ColorService is synced to the URL service
+ */
+ export interface ColorObservedByUrlService {
+  selectedColorOptionName: string;
+  setUrlConfiguration: (urlConfiguration: UrlConfiguration) => void;
+ }
+
 const SELECTED_TAB_UPPER_KEY = 'upper_tab';
 const SELECTED_TAB_LOWER_KEY = 'tab';
 const SELECTED_DATA_KEY = 'selection';
@@ -101,6 +110,7 @@ const LAYOUT_KEY = 'layout';
 const DATA_FIELDS_KEY_SUBSTRING = 'data';
 /** Path to load a new dataset from, on pageload. */
 const NEW_DATASET_PATH = 'new_dataset_path';
+const COLOR_BY_KEY = 'color_by';
 
 const MAX_IDS_IN_URL_SELECTION = 100;
 
@@ -193,6 +203,8 @@ export class UrlService extends LitService {
               `Warning, data index ${dataIndex} is set more than once.`);
         }
         urlConfiguration.dataFields[dataIndex][fieldKey] = value;
+      } else if (key === COLOR_BY_KEY) {
+        urlConfiguration.colorBy = this.urlParseString(value);
       }
     }
     return urlConfiguration;
@@ -231,10 +243,12 @@ export class UrlService extends LitService {
       appState: StateObservedByUrlService,
       modulesService: ModulesObservedByUrlService,
       selectionService: SelectionObservedByUrlService,
-      pinnedSelectionService: SelectionObservedByUrlService) {
+      pinnedSelectionService: SelectionObservedByUrlService,
+      colorService: ColorObservedByUrlService) {
     const urlConfiguration = this.getConfigurationFromUrl();
     appState.setUrlConfiguration(urlConfiguration);
     modulesService.setUrlConfiguration(urlConfiguration);
+    colorService.setUrlConfiguration(urlConfiguration);
 
     const urlSelectedIds = urlConfiguration.selectedData || [];
     selectionService.selectIds(urlSelectedIds, this);
@@ -291,6 +305,9 @@ export class UrlService extends LitService {
       this.setUrlParam(
           urlParams, SELECTED_TAB_LOWER_KEY, modulesService.selectedTabLower);
 
+      this.setUrlParam(urlParams, COLOR_BY_KEY,
+          colorService.selectedColorOptionName);
+
       if (urlParams.toString() !== '') {
         const newUrl = `${window.location.pathname}?${urlParams.toString()}`;
         window.history.replaceState({}, '', newUrl);