ColumnSelect: display column icon when selected

src/components/TableToolbar/ColumnSelect.tsx

@@ -1,7 +1,7 @@
 import { useAtom } from "jotai";
 
 import MultiSelect, { MultiSelectProps } from "@rowy/multiselect";
-import { Stack, Typography } from "@mui/material";
+import { Stack, StackProps, Typography } from "@mui/material";
 
 import { globalScope, altPressAtom } from "@src/atoms/globalScope";
 import { tableScope, tableColumnsOrderedAtom } from "@src/atoms/tableScope";
@@ -30,7 +30,7 @@ export default function ColumnSelect({
     filterColumns
       ? tableColumnsOrdered.filter(filterColumns)
       : tableColumnsOrdered
-  ).map(({ key, name, type, index, fixed }) => ({
+  ).map(({ key, name, type, index }) => ({
     value: key,
     label: name,
     type,
@@ -44,14 +44,39 @@ export default function ColumnSelect({
       labelPlural="columns"
       {...(props as any)}
       itemRenderer={(option: ColumnOption) => <ColumnItem option={option} />}
+      TextFieldProps={{
+        ...props.TextFieldProps,
+        SelectProps: {
+          ...props.TextFieldProps?.SelectProps,
+          renderValue: () => {
+            if (Array.isArray(props.value) && props.value.length > 1)
+              return `${props.value.length} columns`;
+
+            const value = Array.isArray(props.value)
+              ? props.value[0]
+              : props.value;
+            const option = options.find((o) => o.value === value);
+            return option ? (
+              <ColumnItem
+                option={option}
+                sx={{ "& .MuiSvgIcon-root": { my: -0.25 } }}
+              />
+            ) : (
+              value
+            );
+          },
+        },
+      }}
     />
   );
 }
 
-export function ColumnItem({
-  option,
-  children,
-}: React.PropsWithChildren<{ option: ColumnOption }>) {
+export interface IColumnItemProps extends Partial<StackProps> {
+  option: ColumnOption;
+  children?: React.ReactNode;
+}
+
+export function ColumnItem({ option, children, ...props }: IColumnItemProps) {
   const [altPress] = useAtom(altPressAtom, globalScope);
 
   return (
@@ -59,7 +84,11 @@ export function ColumnItem({
       direction="row"
       alignItems="center"
       gap={1}
-      sx={{ color: "text.secondary", width: "100%" }}
+      {...props}
+      sx={[
+        { color: "text.secondary", width: "100%" },
+        ...(Array.isArray(props.sx) ? props.sx : props.sx ? [props.sx] : []),
+      ]}
     >
       {getFieldProp("icon", option.type)}
       <Typography color="text.primary" style={{ flexGrow: 1 }}>