Refactor TreeGrid to have typed items and data

* Both items and data can be typed separately, though it mostly only makes sense to customize one or the other. EntityGroup criteria have generic data and typed items to store extra associated info such as the entity group but the study pages just have typed data because it's simpler. * Pass the item to rowCustomization callbacks since it contains the data as well. Type it appropriately. * Split the root children out so there's no need for an awkward root object.
DataBiosphere · Mar 4, 2025 · ee96d63 · ee96d63
1 parent 4f8ea23
commit ee96d63
Show file tree

Hide file tree

Showing 17 changed files with 225 additions and 268 deletions.
diff --git a/ui/cypress/e2e/featureSet.cy.js b/ui/cypress/e2e/featureSet.cy.js
@@ -7,15 +7,15 @@ describe("Basic tests", () => {
     cy.wait(2000);
 
     cy.iframe().find("a:Contains(Add a data feature)").first().click();
-    cy.iframe().find("[data-testid='tanagra-procedures']").click();
-    cy.iframe().find("input").type("Screening procedure");
-    cy.possiblyMultiSelect("Screening procedure");
-
-    cy.iframe().find("button:Contains(Add data feature)").click();
     cy.iframe().find("[data-testid='tanagra-conditions']").click();
     cy.iframe().find("input").type("Red color");
     cy.possiblyMultiSelect("Red color");
 
+    cy.iframe().find("button:Contains(Add data feature)").click();
+    cy.iframe().find("[data-testid='tanagra-procedures']").click();
+    cy.iframe().find("input").type("Screening procedure");
+    cy.possiblyMultiSelect("Screening procedure");
+
     cy.iframe().find("button:Contains(procedureOccurrence)").click();
     cy.iframe().find("[name='procedure']").click();
 

diff --git a/ui/src/addByCode.tsx b/ui/src/addByCode.tsx
@@ -3,7 +3,12 @@ import Chip from "@mui/material/Chip";
 import { createCriteria, lookupCriteria, LookupEntry } from "cohort";
 import Checkbox from "components/checkbox";
 import Loading from "components/loading";
-import { TreeGrid, TreeGridColumn, TreeGridId } from "components/treeGrid";
+import {
+  TreeGrid,
+  TreeGridColumn,
+  TreeGridId,
+  TreeGridItem,
+} from "components/treeGrid";
 import { useArrayAsTreeGridData } from "components/treeGridHelpers";
 import ActionBar from "actionBar";
 import { DataKey } from "data/types";
@@ -20,7 +25,7 @@ import { cohortURL, useIsSecondBlock } from "router";
 import { insertCohortCriteria, useCohortContext } from "cohortContext";
 import { isValid } from "util/valid";
 
-type LookupEntryItem = {
+type LookupEntryData = {
   config?: JSX.Element;
   code: DataKey;
   name?: string;
@@ -44,7 +49,7 @@ export function AddByCode() {
     [underlay.criteriaSelectors]
   );
 
-  const lookupEntriesState = useSWRMutation<LookupEntryItem[]>(
+  const lookupEntriesState = useSWRMutation<LookupEntryData[]>(
     {
       component: "AddByCode",
       query,
@@ -106,7 +111,7 @@ export function AddByCode() {
   const data = useArrayAsTreeGridData(lookupEntriesState?.data ?? [], "code");
 
   const onInsert = useCallback(() => {
-    const configMap = new Map<string, LookupEntryItem[]>();
+    const configMap = new Map<string, LookupEntryData[]>();
     lookupEntriesState.data?.forEach((e) => {
       if (!e.entry || !selected.has(e.code)) {
         return;
@@ -203,19 +208,13 @@ export function AddByCode() {
         </GridBox>
         <Loading immediate showProgressOnMutate status={lookupEntriesState}>
           {lookupEntriesState.data?.length ? (
-            <TreeGrid
+            <TreeGrid<TreeGridItem<LookupEntryData>>
               columns={columns}
               data={data}
-              rowCustomization={(id: TreeGridId) => {
-                if (!lookupEntriesState.data) {
-                  return undefined;
-                }
-
-                const item = data.get(id)?.data as LookupEntryItem;
-                if (!item) {
-                  return undefined;
-                }
-
+              rowCustomization={(
+                id: TreeGridId,
+                { data }: TreeGridItem<LookupEntryData>
+              ) => {
                 const sel = selected.has(id);
                 return [
                   {
@@ -225,7 +224,7 @@ export function AddByCode() {
                         size="small"
                         fontSize="inherit"
                         checked={sel}
-                        disabled={!item.entry}
+                        disabled={!data.entry}
                         onChange={() => {
                           updateSelected((selected) => {
                             if (sel) {

diff --git a/ui/src/addCohort.tsx b/ui/src/addCohort.tsx
@@ -48,7 +48,7 @@ export function AddCohort() {
   const navigate = useNavigate();
   const params = useBaseParams();
 
-  const cohortsState = useSWR(
+  const cohortsState = useSWR<CohortData[]>(
     {
       type: "cohorts",
       studyId,
@@ -109,16 +109,7 @@ export function AddCohort() {
             <TreeGrid
               data={data}
               columns={columns}
-              rowCustomization={(id: TreeGridId) => {
-                if (!cohortsState.data) {
-                  return undefined;
-                }
-
-                const cohortData = data.get(id)?.data as CohortData;
-                if (!cohortData) {
-                  return undefined;
-                }
-
+              rowCustomization={(id: TreeGridId, { data: cohortData }) => {
                 return [
                   {
                     column: columns.length - 2,

diff --git a/ui/src/addCriteria.tsx b/ui/src/addCriteria.tsx
@@ -23,7 +23,7 @@ import {
   TreeGrid,
   TreeGridId,
   TreeGridItem,
-  TreeGridRowData,
+  TreeGridData,
 } from "components/treeGrid";
 import { MergedItem } from "data/mergeLists";
 import { Criteria, isTemporalSection } from "data/source";
@@ -366,9 +366,7 @@ function AddCriteria(props: AddCriteriaProps) {
 
   const search = useCallback(async () => {
     const children: DataKey[] = [];
-    const data = new Map<TreeGridId, CriteriaItem>([
-      ["root", { data: {}, children }],
-    ]);
+    const rows = new Map<TreeGridId, CriteriaItem>();
 
     if (query) {
       const res = await searchCriteria(
@@ -390,13 +388,16 @@ function AddCriteria(props: AddCriteriaProps) {
           },
           entry: entry,
         };
-        data.set(key, item);
+        rows.set(key, item);
       });
     }
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [underlaySource, query, selectedOptions, optionsMap]);
-  const searchState = useSWRImmutable<Map<TreeGridId, CriteriaItem>>(
+  const searchState = useSWRImmutable<TreeGridData<CriteriaItem>>(
     {
       component: "AddCriteria",
       underlayName: underlay.name,
@@ -540,7 +541,7 @@ function AddCriteria(props: AddCriteriaProps) {
         {query ? (
           <Paper>
             <Loading status={searchState}>
-              {!searchState.data?.get("root")?.children?.length ? (
+              {!searchState.data?.children?.length ? (
                 <Empty
                   minHeight="300px"
                   image={emptyImage}
@@ -549,17 +550,9 @@ function AddCriteria(props: AddCriteriaProps) {
               ) : (
                 <TreeGrid
                   columns={columns}
-                  data={searchState.data ?? new Map()}
-                  rowCustomization={(
-                    id: TreeGridId,
-                    rowData: TreeGridRowData
-                  ) => {
-                    if (!searchState.data) {
-                      return undefined;
-                    }
-
-                    const item = searchState.data.get(id);
-                    const option = optionsMap.get(item?.entry?.source ?? "");
+                  data={searchState.data}
+                  rowCustomization={(id, item) => {
+                    const option = optionsMap.get(item.entry?.source ?? "");
                     if (!option || !item?.entry?.source) {
                       throw new Error(
                         `Item source "${item?.entry?.source}" doesn't match any criteria config ID.`
@@ -572,7 +565,9 @@ function AddCriteria(props: AddCriteriaProps) {
                         content: (
                           <GridLayout colAlign="center">
                             <Button
-                              data-testid={rowData[searchConfig.columns[0].key]}
+                              data-testid={
+                                item.data[searchConfig.columns[0].key]
+                              }
                               onClick={() => onClick(option, item.entry?.data)}
                               variant="outlined"
                             >

diff --git a/ui/src/cohortReview/cohortReviewList.tsx b/ui/src/cohortReview/cohortReviewList.tsx
@@ -679,16 +679,7 @@ function Annotations() {
             <TreeGrid
               data={data}
               columns={columns}
-              rowCustomization={(id: TreeGridId) => {
-                if (!annotationsState.data) {
-                  return undefined;
-                }
-
-                const annotation = data.get(id)?.data;
-                if (!annotation) {
-                  return undefined;
-                }
-
+              rowCustomization={(id: TreeGridId, { data: annotation }) => {
                 return [
                   {
                     column: columns.length - 1,

diff --git a/ui/src/cohortReview/participantsListDialog.tsx b/ui/src/cohortReview/participantsListDialog.tsx
@@ -11,7 +11,7 @@ import {
   useReviewSearchState,
 } from "cohortReview/reviewHooks";
 import Loading from "components/loading";
-import { TreeGrid, TreeGridData, TreeGridId } from "components/treeGrid";
+import { TreeGrid, TreeGridId } from "components/treeGrid";
 import GridLayout from "layout/gridLayout";
 import React, { useMemo } from "react";
 
@@ -108,7 +108,8 @@ function ParticipantsList(props: ParticipantsListProps) {
   );
 
   const data = useMemo(() => {
-    const data: TreeGridData = new Map([["root", { data: {}, children: [] }]]);
+    const children: TreeGridId[] = [];
+    const rows = new Map();
 
     instancesState.data
       ?.slice(
@@ -117,21 +118,24 @@ function ParticipantsList(props: ParticipantsListProps) {
       )
       .forEach((instance) => {
         const key = instance.data.key;
-        data.set(key, { data: { ...instance.data } });
-        data.get("root")?.children?.push(key);
+        rows.set(key, { data: { ...instance.data } });
+        children.push(key);
 
         annotationsState.data?.forEach((a) => {
           const values = instance.annotations.get(a.id);
           if (values) {
-            const valueData = data.get(key)?.data;
+            const valueData = rows.get(key)?.data;
             if (valueData) {
               valueData[`t_${a.id}`] = values[values.length - 1].value;
             }
           }
         });
       });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [instancesState, annotationsState, props.page, props.rowsPerPage]);
 
   return (

diff --git a/ui/src/cohortReview/plugins/occurrenceTable.tsx b/ui/src/cohortReview/plugins/occurrenceTable.tsx
@@ -6,7 +6,6 @@ import {
 import {
   TreeGrid,
   TreeGridColumn,
-  TreeGridData,
   TreeGridSortOrder,
 } from "components/treeGrid";
 import { TreeGridSortDirection } from "components/treeGridHelpers";
@@ -63,14 +62,17 @@ export function OccurrenceTable({
 
   const data = useMemo(() => {
     const children: DataKey[] = [];
-    const data: TreeGridData = new Map([["root", { data: {}, children }]]);
+    const rows = new Map();
 
     context.rows[config.entity]?.forEach((o) => {
-      data.set(o.key, { data: o });
+      rows.set(o.key, { data: o });
       children.push(o.key);
     });
 
-    return data;
+    return {
+      rows,
+      children,
+    };
   }, [context, config]);
 
   const filterRegExps = useMemo(() => {
@@ -83,25 +85,24 @@ export function OccurrenceTable({
 
   const sortedData = useMemo(() => {
     return produce(data, (data) => {
-      const root = data.get("root");
-      if (!root) {
-        return;
-      }
-
-      root.children = (root.children ?? []).filter((child) =>
+      data.children = (data.children ?? []).filter((child) =>
         Object.entries(filterRegExps ?? {}).reduce(
           (cur: boolean, [col, re]) =>
             cur &&
             re.test(
-              stringifyDataValue(data.get(child)?.data?.[col] as DataValue)
+              stringifyDataValue(data.rows.get(child)?.data?.[col] as DataValue)
             ),
           true
         )
       );
-      root.children.sort((a, b) => {
+      data.children.sort((a, b) => {
         for (const o of searchState.sortOrders ?? []) {
-          const valA = data.get(a)?.data?.[o.column] as DataValue | undefined;
-          const valB = data.get(b)?.data?.[o.column] as DataValue | undefined;
+          const valA = data.rows.get(a)?.data?.[o.column] as
+            | DataValue
+            | undefined;
+          const valB = data.rows.get(b)?.data?.[o.column] as
+            | DataValue
+            | undefined;
           const c = compareDataValues(valA, valB);
           if (c !== 0) {
             return o.direction === TreeGridSortDirection.Asc ? c : -c;