Add ability to highlight/select method rows

Author: shati-patel

extensions/ql-vscode/src/view/common/DataGrid.tsx

@@ -37,7 +37,10 @@ export function DataGrid({ gridTemplateColumns, children }: DataGridProps) {
   );
 }
 
-const StyledDataGridRow = styled.div<{ $focused?: boolean }>`
+const StyledDataGridRow = styled.div<{
+  $focused?: boolean;
+  $selected?: boolean;
+}>`
   display: contents;
 
   &:hover > * {
@@ -47,15 +50,20 @@ const StyledDataGridRow = styled.div<{ $focused?: boolean }>`
   & > * {
     // Use !important to override the background color set by the hover state
     background-color: ${(props) =>
+      // TODO: Use different colors for selected and focused
       props.$focused
         ? "var(--vscode-editor-selectionBackground) !important"
-        : "inherit"};
+        : props.$selected
+          ? "var(--vscode-editor-selectionBackground) !important"
+          : "inherit"};
   }
 `;
 
 interface DataGridRowProps {
   focused?: boolean;
+  selected?: boolean;
   children: ReactNode;
+  onClick?: () => void;
   "data-testid"?: string;
 }
 
@@ -69,10 +77,22 @@ interface DataGridRowProps {
  */
 export const DataGridRow = forwardRef(
   (
-    { focused, children, "data-testid": testId }: DataGridRowProps,
+    {
+      focused,
+      selected,
+      children,
+      "data-testid": testId,
+      onClick,
+    }: DataGridRowProps,
     ref?: React.Ref<HTMLElement | undefined>,
   ) => (
-    <StyledDataGridRow $focused={focused} ref={ref} data-testid={testId}>
+    <StyledDataGridRow
+      $focused={focused}
+      $selected={selected}
+      ref={ref}
+      data-testid={testId}
+      onClick={onClick}
+    >
       {children}
     </StyledDataGridRow>
   ),

extensions/ql-vscode/src/view/model-editor/LibraryRow.tsx

@@ -72,11 +72,13 @@ export type LibraryRowProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
+  selectedSignatures: Set<string>;
   inProgressMethods: Set<string>;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
   onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
+  onMethodClick: (methodSignature: string) => void;
   onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     dependencyName: string,
@@ -93,11 +95,13 @@ export const LibraryRow = ({
   methods,
   modeledMethodsMap,
   modifiedSignatures,
+  selectedSignatures,
   inProgressMethods,
   viewState,
   hideModeledMethods,
   revealedMethodSignature,
   onChange,
+  onMethodClick,
   onSaveModelClick,
   onGenerateFromLlmClick,
   onStopGenerateFromLlmClick,
@@ -229,11 +233,13 @@ export const LibraryRow = ({
             methods={methods}
             modeledMethodsMap={modeledMethodsMap}
             modifiedSignatures={modifiedSignatures}
+            selectedSignatures={selectedSignatures}
             inProgressMethods={inProgressMethods}
             viewState={viewState}
             hideModeledMethods={hideModeledMethods}
             revealedMethodSignature={revealedMethodSignature}
             onChange={onChange}
+            onMethodClick={onMethodClick}
           />
           <SectionDivider />
           <ButtonsContainer>

extensions/ql-vscode/src/view/model-editor/MethodRow.tsx

@@ -71,10 +71,12 @@ export type MethodRowProps = {
   methodCanBeModeled: boolean;
   modeledMethods: ModeledMethod[];
   methodIsUnsaved: boolean;
+  methodIsSelected: boolean;
   modelingInProgress: boolean;
   viewState: ModelEditorViewState;
   revealedMethodSignature: string | null;
   onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
+  onMethodClick: (methodSignature: string) => void;
 };
 
 export const MethodRow = (props: MethodRowProps) => {
@@ -104,9 +106,11 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       method,
       modeledMethods: modeledMethodsProp,
       methodIsUnsaved,
+      methodIsSelected,
       viewState,
       revealedMethodSignature,
       onChange,
+      onMethodClick,
     } = props;
 
     const [focusedIndex, setFocusedIndex] = useState<number | null>(null);
@@ -187,6 +191,10 @@ const ModelableMethodRow = forwardRef<HTMLElement | undefined, MethodRowProps>(
       <DataGridRow
         data-testid="modelable-method-row"
         focused={revealedMethodSignature === method.signature}
+        selected={methodIsSelected}
+        onClick={() => {
+          onMethodClick(method.signature);
+        }}
       >
         <DataGridCell
           gridRow={`span ${modeledMethods.length + validationErrors.length}`}

extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx

@@ -96,6 +96,10 @@ export function ModelEditor({
     new Set(),
   );
 
+  const [selectedSignatures, setSelectedSignatures] = useState<Set<string>>(
+    new Set(),
+  );
+
   const [inProgressMethods, setInProgressMethods] = useState<Set<string>>(
     new Set(),
   );
@@ -190,6 +194,19 @@ export function ModelEditor({
     [],
   );
 
+  const onMethodClick = useCallback(
+    (methodSignature: string) => {
+      const newSelectedSignatures = new Set(selectedSignatures);
+      if (selectedSignatures.has(methodSignature)) {
+        newSelectedSignatures.delete(methodSignature);
+      } else {
+        newSelectedSignatures.add(methodSignature);
+      }
+      setSelectedSignatures(newSelectedSignatures);
+    },
+    [selectedSignatures],
+  );
+
   const onRefreshClick = useCallback(() => {
     vscode.postMessage({
       t: "refreshMethods",
@@ -340,11 +357,13 @@ export function ModelEditor({
           methods={methods}
           modeledMethodsMap={modeledMethods}
           modifiedSignatures={modifiedSignatures}
+          selectedSignatures={selectedSignatures}
           inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledMethods={hideModeledMethods}
           revealedMethodSignature={revealedMethodSignature}
           onChange={onChange}
+          onMethodClick={onMethodClick}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
           onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}

extensions/ql-vscode/src/view/model-editor/ModeledMethodDataGrid.tsx

@@ -16,22 +16,26 @@ export type ModeledMethodDataGridProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
+  selectedSignatures: Set<string>;
   inProgressMethods: Set<string>;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
   onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
+  onMethodClick: (methodSignature: string) => void;
 };
 
 export const ModeledMethodDataGrid = ({
   methods,
   modeledMethodsMap,
   modifiedSignatures,
+  selectedSignatures,
   inProgressMethods,
   viewState,
   hideModeledMethods,
   revealedMethodSignature,
   onChange,
+  onMethodClick,
 }: ModeledMethodDataGridProps) => {
   const [methodsWithModelability, numHiddenMethods]: [
     Array<{ method: Method; methodCanBeModeled: boolean }>,
@@ -80,10 +84,12 @@ export const ModeledMethodDataGrid = ({
                 methodCanBeModeled={methodCanBeModeled}
                 modeledMethods={modeledMethods}
                 methodIsUnsaved={modifiedSignatures.has(method.signature)}
+                methodIsSelected={selectedSignatures.has(method.signature)}
                 modelingInProgress={inProgressMethods.has(method.signature)}
                 viewState={viewState}
                 revealedMethodSignature={revealedMethodSignature}
                 onChange={onChange}
+                onMethodClick={onMethodClick}
               />
             );
           })}

extensions/ql-vscode/src/view/model-editor/ModeledMethodsList.tsx

@@ -14,11 +14,13 @@ export type ModeledMethodsListProps = {
   methods: Method[];
   modeledMethodsMap: Record<string, ModeledMethod[]>;
   modifiedSignatures: Set<string>;
+  selectedSignatures: Set<string>;
   inProgressMethods: Set<string>;
   revealedMethodSignature: string | null;
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   onChange: (methodSignature: string, modeledMethods: ModeledMethod[]) => void;
+  onMethodClick: (methodSignature: string) => void;
   onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     packageName: string,
@@ -37,11 +39,13 @@ export const ModeledMethodsList = ({
   methods,
   modeledMethodsMap,
   modifiedSignatures,
+  selectedSignatures,
   inProgressMethods,
   viewState,
   hideModeledMethods,
   revealedMethodSignature,
   onChange,
+  onMethodClick,
   onSaveModelClick,
   onGenerateFromLlmClick,
   onStopGenerateFromLlmClick,
@@ -83,11 +87,13 @@ export const ModeledMethodsList = ({
           methods={grouped[libraryName]}
           modeledMethodsMap={modeledMethodsMap}
           modifiedSignatures={modifiedSignatures}
+          selectedSignatures={selectedSignatures}
           inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledMethods={hideModeledMethods}
           revealedMethodSignature={revealedMethodSignature}
           onChange={onChange}
+          onMethodClick={onMethodClick}
           onSaveModelClick={onSaveModelClick}
           onGenerateFromLlmClick={onGenerateFromLlmClick}
           onStopGenerateFromLlmClick={onStopGenerateFromLlmClick}