From 91d6b683f757660443725e715c219db3de1372f5 Mon Sep 17 00:00:00 2001
From: Peyton Lee <peyton.lee@alleninstitute.org>
Date: Mon, 6 Jan 2025 11:57:45 -0800
Subject: [PATCH] refactor: Replace toggleLabelOnId

---
 src/colorizer/AnnotationData.ts    | 43 +++++++++++++++---------------
 src/colorizer/utils/react_utils.ts |  8 +++---
 2 files changed, 25 insertions(+), 26 deletions(-)

diff --git a/src/colorizer/AnnotationData.ts b/src/colorizer/AnnotationData.ts
index 5a59b75d..73045c0a 100644
--- a/src/colorizer/AnnotationData.ts
+++ b/src/colorizer/AnnotationData.ts
@@ -30,6 +30,15 @@ export interface IAnnotationDataGetters {
    */
   getLabelsAppliedToId(id: number): number[];
 
+  /**
+   * Returns whether the label has been applied to the object ID.
+   * @param labelIdx The index of the label to look up.
+   * @param id The object ID.
+   * @returns `true` if the label has been applied to the object ID, `false`
+   * otherwise.
+   */
+  getLabelOnId(labelIdx: number, id: number): boolean;
+
   /**
    * Returns all object IDs that the label has been applied to.
    * @param labelIdx The index of the label to look up object IDs for.
@@ -83,9 +92,7 @@ export interface IAnnotationDataSetters {
   setLabelColor(labelIdx: number, color: Color): void;
   deleteLabel(labelIdx: number): void;
 
-  applyLabelToId(labelIdx: number, id: number): void;
-  removeLabelFromId(labelIdx: number, id: number): void;
-  toggleLabelOnId(labelIdx: number, id: number): void;
+  setLabelOnId(labelIdx: number, id: number, value: boolean): void;
 }
 
 export type IAnnotationData = IAnnotationDataGetters & IAnnotationDataSetters;
@@ -113,9 +120,8 @@ export class AnnotationData implements IAnnotationDataGetters, IAnnotationDataSe
     this.setLabelName = this.setLabelName.bind(this);
     this.setLabelColor = this.setLabelColor.bind(this);
     this.deleteLabel = this.deleteLabel.bind(this);
-    this.toggleLabelOnId = this.toggleLabelOnId.bind(this);
-    this.applyLabelToId = this.applyLabelToId.bind(this);
-    this.removeLabelFromId = this.removeLabelFromId.bind(this);
+    this.getLabelOnId = this.getLabelOnId.bind(this);
+    this.setLabelOnId = this.setLabelOnId.bind(this);
   }
 
   // Getters
@@ -124,6 +130,11 @@ export class AnnotationData implements IAnnotationDataGetters, IAnnotationDataSe
     return [...this.labelData];
   }
 
+  getLabelOnId(labelIdx: number, id: number): boolean {
+    this.validateIndex(labelIdx);
+    return this.labelData[labelIdx].ids.has(id);
+  }
+
   getLabelsAppliedToId(id: number): number[] {
     const labelIdxs: number[] = [];
     for (let i = 0; i < this.labelData.length; i++) {
@@ -210,23 +221,13 @@ export class AnnotationData implements IAnnotationDataGetters, IAnnotationDataSe
     this.timeToLabelIdMap = null;
   }
 
-  applyLabelToId(labelIdx: number, id: number): void {
-    this.validateIndex(labelIdx);
-    this.labelData[labelIdx].ids.add(id);
-    this.timeToLabelIdMap = null;
-  }
-
-  removeLabelFromId(labelIdx: number, id: number): void {
+  setLabelOnId(labelIdx: number, id: number, value: boolean): void {
     this.validateIndex(labelIdx);
-    this.labelData[labelIdx].ids.delete(id);
-    this.timeToLabelIdMap = null;
-  }
-
-  toggleLabelOnId(labelIdx: number, id: number): void {
-    if (!this.labelData[labelIdx].ids.has(id)) {
-      this.applyLabelToId(labelIdx, id);
+    if (value) {
+      this.labelData[labelIdx].ids.add(id);
     } else {
-      this.removeLabelFromId(labelIdx, id);
+      this.labelData[labelIdx].ids.delete(id);
     }
+    this.timeToLabelIdMap = null;
   }
 }
diff --git a/src/colorizer/utils/react_utils.ts b/src/colorizer/utils/react_utils.ts
index ff3cecd8..f7d7cce2 100644
--- a/src/colorizer/utils/react_utils.ts
+++ b/src/colorizer/utils/react_utils.ts
@@ -372,7 +372,7 @@ export const useAnnotations = (): AnnotationState => {
     } else if (currentLabelIdx === labelIdx) {
       setCurrentLabelIdx(null);
       setIsAnnotationEnabled(false);
-    }else if (currentLabelIdx > labelIdx) {
+    } else if (currentLabelIdx > labelIdx) {
       // Decrement because all indices will shift over
       setCurrentLabelIdx(currentLabelIdx - 1);
     }
@@ -385,6 +385,7 @@ export const useAnnotations = (): AnnotationState => {
       getLabelsAppliedToId: annotationData.getLabelsAppliedToId,
       getLabeledIds: annotationData.getLabeledIds,
       getTimeToLabelIdMap: annotationData.getTimeToLabelIdMap,
+      getLabelOnId: annotationData.getLabelOnId,
     })
   , [dataUpdateCounter]);
 
@@ -403,9 +404,6 @@ export const useAnnotations = (): AnnotationState => {
     setLabelName: wrapFunctionInUpdate(annotationData.setLabelName),
     setLabelColor: wrapFunctionInUpdate(annotationData.setLabelColor),
     deleteLabel: wrapFunctionInUpdate(onDeleteLabel),
-
-    applyLabelToId: wrapFunctionInUpdate(annotationData.applyLabelToId),
-    removeLabelFromId: wrapFunctionInUpdate(annotationData.removeLabelFromId),
-    toggleLabelOnId: wrapFunctionInUpdate(annotationData.toggleLabelOnId),
+    setLabelOnId: wrapFunctionInUpdate(annotationData.setLabelOnId),
   };
 };