Select: Support option groups in virtualised Selects (#88232)

* change select group headers to always be visible

* remove unnecessary SelectOptionGroup

* hacky change to work with virtualised lists

* undo this change

* add top border

* remove children from the category header

* don't virtualize story

* bit of renaming

* comments

* add new story for virtualized options

packages/grafana-ui/package.json

@@ -132,6 +132,7 @@
     "@testing-library/jest-dom": "6.4.2",
     "@testing-library/react": "15.0.2",
     "@testing-library/user-event": "14.5.2",
+    "@types/chance": "1.1.6",
     "@types/common-tags": "^1.8.0",
     "@types/d3": "7.4.3",
     "@types/hoist-non-react-statics": "3.3.5",
@@ -158,6 +159,7 @@
     "@types/testing-library__jest-dom": "5.14.9",
     "@types/tinycolor2": "1.4.6",
     "@types/uuid": "9.0.8",
+    "chance": "1.1.11",
     "common-tags": "1.8.2",
     "core-js": "3.37.0",
     "css-loader": "7.1.1",

packages/grafana-ui/src/components/Select/Select.story.tsx

@@ -1,6 +1,7 @@
 import { auto } from '@popperjs/core';
 import { action } from '@storybook/addon-actions';
 import { Meta, StoryFn } from '@storybook/react';
+import Chance from 'chance';
 import React, { useState } from 'react';
 
 import { SelectableValue, toIconName } from '@grafana/data';
@@ -12,6 +13,26 @@ import mdx from './Select.mdx';
 import { generateOptions, generateThousandsOfOptions } from './mockOptions';
 import { SelectCommonProps } from './types';
 
+const chance = new Chance();
+
+const manyGroupedOptions = [
+  { label: 'Foo', value: '1' },
+  {
+    label: 'Animals',
+    options: new Array(100).fill(0).map((_, i) => {
+      const animal = chance.animal();
+      return { label: animal, value: animal };
+    }),
+  },
+  {
+    label: 'People',
+    options: new Array(100).fill(0).map((_, i) => {
+      const person = chance.name();
+      return { label: person, value: person };
+    }),
+  },
+];
+
 const meta: Meta = {
   title: 'Forms/Select',
   component: Select,
@@ -242,6 +263,26 @@ export const MultiSelectWithOptionGroups: StoryFn = (args) => {
   );
 };
 
+export const MultiSelectWithOptionGroupsVirtualized: StoryFn = (args) => {
+  const [value, setValue] = useState<string[]>();
+
+  return (
+    <>
+      <MultiSelect
+        options={manyGroupedOptions}
+        virtualized
+        value={value}
+        onChange={(v) => {
+          setValue(v.map((v) => v.value!));
+          action('onChange')(v);
+        }}
+        prefix={getPrefix(args.icon)}
+        {...args}
+      />
+    </>
+  );
+};
+
 export const MultiSelectBasic: StoryFn = (args) => {
   const [value, setValue] = useState<Array<SelectableValue<string>>>([]);
 

packages/grafana-ui/src/components/Select/SelectMenu.tsx

@@ -1,6 +1,6 @@
 import { cx } from '@emotion/css';
 import { max } from 'lodash';
-import React, { RefCallback, useEffect, useRef } from 'react';
+import React, { RefCallback, useEffect, useMemo, useRef } from 'react';
 import { MenuListProps } from 'react-select';
 import { FixedSizeList as List } from 'react-window';
 
@@ -57,8 +57,18 @@ export const VirtualizedSelectMenu = ({
   const styles = getSelectStyles(theme);
   const listRef = useRef<List>(null);
 
-  const focusedIndex = options.findIndex((option: SelectableValue<unknown>) => option.value === focusedOption?.value);
+  // we need to check for option groups (categories)
+  // these are top level options with child options
+  // if they exist, flatten the list of options
+  const flattenedOptions = useMemo(
+    () => options.flatMap((option) => (option.options ? [option, ...option.options] : [option])),
+    [options]
+  );
 
+  // scroll the focused option into view when navigating with keyboard
+  const focusedIndex = flattenedOptions.findIndex(
+    (option: SelectableValue<unknown>) => option.value === focusedOption?.value
+  );
   useEffect(() => {
     listRef.current?.scrollToItem(focusedIndex);
   }, [focusedIndex]);
@@ -67,10 +77,23 @@ export const VirtualizedSelectMenu = ({
     return null;
   }
 
-  const longestOption = max(options.map((option) => option.label?.length)) ?? 0;
+  // flatten the children to account for any categories
+  // these will have array children that are the individual options
+  const flattenedChildren = children.flatMap((child) => {
+    if (hasArrayChildren(child)) {
+      // need to remove the children from the category else they end up in the DOM twice
+      const childWithoutChildren = React.cloneElement(child, {
+        children: null,
+      });
+      return [childWithoutChildren, ...child.props.children];
+    }
+    return [child];
+  });
+
+  const longestOption = max(flattenedOptions.map((option) => option.label?.length)) ?? 0;
   const widthEstimate =
     longestOption * VIRTUAL_LIST_WIDTH_ESTIMATE_MULTIPLIER + VIRTUAL_LIST_PADDING * 2 + VIRTUAL_LIST_WIDTH_EXTRA;
-  const heightEstimate = Math.min(options.length * VIRTUAL_LIST_ITEM_HEIGHT, maxHeight);
+  const heightEstimate = Math.min(flattenedChildren.length * VIRTUAL_LIST_ITEM_HEIGHT, maxHeight);
 
   return (
     <List
@@ -79,14 +102,20 @@ export const VirtualizedSelectMenu = ({
       height={heightEstimate}
       width={widthEstimate}
       aria-label="Select options menu"
-      itemCount={children.length}
+      itemCount={flattenedChildren.length}
       itemSize={VIRTUAL_LIST_ITEM_HEIGHT}
     >
-      {({ index, style }) => <div style={{ ...style, overflow: 'hidden' }}>{children[index]}</div>}
+      {({ index, style }) => <div style={{ ...style, overflow: 'hidden' }}>{flattenedChildren[index]}</div>}
     </List>
   );
 };
 
+// check if a child has array children (and is therefore a react-select group)
+// we need to flatten these so the correct count and elements are passed to the virtualized list
+const hasArrayChildren = (child: React.ReactNode) => {
+  return React.isValidElement(child) && Array.isArray(child.props.children);
+};
+
 VirtualizedSelectMenu.displayName = 'VirtualizedSelectMenu';
 
 interface SelectMenuOptionProps<T> {

yarn.lock

@@ -3648,6 +3648,7 @@ __metadata:
     "@testing-library/jest-dom": "npm:6.4.2"
     "@testing-library/react": "npm:15.0.2"
     "@testing-library/user-event": "npm:14.5.2"
+    "@types/chance": "npm:1.1.6"
     "@types/common-tags": "npm:^1.8.0"
     "@types/d3": "npm:7.4.3"
     "@types/hoist-non-react-statics": "npm:3.3.5"
@@ -3676,6 +3677,7 @@ __metadata:
     "@types/uuid": "npm:9.0.8"
     ansicolor: "npm:1.1.100"
     calculate-size: "npm:1.1.1"
+    chance: "npm:1.1.11"
     classnames: "npm:2.5.1"
     common-tags: "npm:1.8.2"
     core-js: "npm:3.37.0"
@@ -7827,7 +7829,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@types/chance@npm:^1.1.3":
+"@types/chance@npm:1.1.6, @types/chance@npm:^1.1.3":
   version: 1.1.6
   resolution: "@types/chance@npm:1.1.6"
   checksum: 10/f4366f1b3144d143af3e6f0fad2ed1db7b9bdfa7d82d40944e9619d57fe7e6b60e8c1452f47a8ededa6b2188932879518628ecd9aac81c40384ded39c26338ba
@@ -11639,7 +11641,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"chance@npm:^1.0.10":
+"chance@npm:1.1.11, chance@npm:^1.0.10":
   version: 1.1.11
   resolution: "chance@npm:1.1.11"
   checksum: 10/d76cc76dbcb837f051e6080d94be8bdcd07559d52077854f614c7081062fee10d4c3b508f6ae4b70303dac000422ea35160b01de165fcd176a01f67ea4b2cef5