Make SegmentedControl support non-string values

packages/ui/src/SegmentedControl/index.test.tsx

@@ -35,4 +35,46 @@ describe('<SegmentedControl />', () => {
 
     expect(onChange).toHaveBeenCalledWith('two');
   });
+
+  describe('when the options have non-string values', () => {
+    const valueOne = { toString: () => 'one' };
+    const valueTwo = { toString: () => 'two' };
+    const valueThree = { toString: () => 'three' };
+
+    const options = [
+      { value: valueOne, label: 'One' },
+      { value: valueTwo, label: 'Two' },
+      { value: valueThree, label: 'Three' },
+    ];
+
+    it('calls the `onClick` handler with the value of the clicked option', () => {
+      const { getByText } = render(
+        <SegmentedControl value={valueOne} options={options} onChange={onChange} />,
+        { wrapper: PenumbraUIProvider },
+      );
+      fireEvent.click(getByText('Two', { selector: ':not([aria-hidden])' }));
+
+      expect(onChange).toHaveBeenCalledWith(valueTwo);
+    });
+
+    describe("when the options' `.toString()` methods return non-unique values", () => {
+      const valueOne = { toString: () => 'one' };
+      const valueTwo = { toString: () => 'two' };
+      const valueTwoAgain = { toString: () => 'two' };
+
+      const options = [
+        { value: valueOne, label: 'One' },
+        { value: valueTwo, label: 'Two' },
+        { value: valueTwoAgain, label: 'Two again' },
+      ];
+
+      it('throws', () => {
+        expect(() =>
+          render(<SegmentedControl value={valueOne} options={options} onChange={onChange} />, {
+            wrapper: PenumbraUIProvider,
+          }),
+        ).toThrow('The value options passed to `<SegmentedControl />` are not unique.');
+      });
+    });
+  });
 });

packages/ui/src/SegmentedControl/index.tsx

@@ -5,6 +5,8 @@ import { Density } from '../types/Density';
 import { useDensity } from '../hooks/useDensity';
 import * as RadixRadioGroup from '@radix-ui/react-radio-group';
 import { useDisabled } from '../hooks/useDisabled';
+import { ToStringable } from '../utils/ToStringable';
+import { useEffect } from 'react';
 
 const Root = styled.div`
   display: flex;
@@ -34,17 +36,44 @@ const Segment = styled.button<{
   padding-right: ${props => props.theme.spacing(props.$density === 'sparse' ? 4 : 2)};
 `;
 
-export interface Option {
-  value: string;
+/**
+ * Radix's `<RadioGroup />` component only accepts strings for its values, but
+ * we don't want to enforce that in `<SegmentedControl />`. Instead, we allow
+ * options to be passed whose values extend `ToStringable` (i.e., they have a
+ * `.toString()` method). Then, when a specific option is selected and passed to
+ * `onChange()`, we need to map from the string value back to the original value
+ * passed in the options array.
+ *
+ * To make sure this works as expected, we need to assert that each option
+ * value's `.toString()` method returns a unique value. That way, we can avoid a
+ * situation where, e.g., all the options' values return `[object Object]`, and
+ * the wrong object is passed to `onChange`.
+ */
+const assertUniqueOptions = (options: Option<ToStringable>[]) => {
+  const existingOptions = new Set<string>();
+
+  options.forEach(option => {
+    if (existingOptions.has(option.value.toString())) {
+      throw new Error(
+        'The value options passed to `<SegmentedControl />` are not unique. Please check that the result of calling `.toString()` on each of the options passed to `<SegmentedControl />` is unique.',
+      );
+    }
+
+    existingOptions.add(option.value.toString());
+  });
+};
+
+export interface Option<ValueType extends ToStringable> {
+  value: ValueType;
   label: string;
   /** Whether this individual option should be disabled. */
   disabled?: boolean;
 }
 
-export interface SegmentedControlProps {
-  value: string;
-  onChange: (value: string) => void;
-  options: Option[];
+export interface SegmentedControlProps<ValueType extends ToStringable> {
+  value: ValueType;
+  onChange: (value: ValueType) => void;
+  options: Option<ValueType>[];
   /**
    * Whether this entire control should be disabled. Note that single options
    * can be disabled individually by setting the `disabled` property for that
@@ -74,15 +103,31 @@ export interface SegmentedControlProps {
  * />
  * ```
  */
-export const SegmentedControl = ({ value, onChange, options, disabled }: SegmentedControlProps) => {
+export const SegmentedControl = <ValueType extends ToStringable>({
+  value,
+  onChange,
+  options,
+  disabled,
+}: SegmentedControlProps<ValueType>) => {
   const density = useDensity();
   disabled = useDisabled(disabled);
 
+  useEffect(() => assertUniqueOptions(options), [options]);
+
+  const handleChange = (value: string) => {
+    const matchingOption = options.find(option => option.value.toString() === value)!;
+    onChange(matchingOption.value);
+  };
+
   return (
-    <RadixRadioGroup.Root asChild value={value} onValueChange={onChange}>
+    <RadixRadioGroup.Root asChild value={value.toString()} onValueChange={handleChange}>
       <Root>
         {options.map(option => (
-          <RadixRadioGroup.Item asChild key={option.value} value={option.value}>
+          <RadixRadioGroup.Item
+            asChild
+            key={option.value.toString()}
+            value={option.value.toString()}
+          >
             <Segment
               onClick={() => onChange(option.value)}
               $getBorderRadius={theme => theme.borderRadius.full}

packages/ui/src/utils/ToStringable.ts

@@ -0,0 +1,9 @@
+/**
+ * Utility interface to represent types that can be cast to string. Useful for
+ * e.g., accepting an array of `.toString()`-able items will be mapped over, so
+ * that the items can have `.toString()` called on them for the React `key`
+ * prop.
+ */
+export interface ToStringable {
+  toString: () => string;
+}