Forward the ref in the Button component

src/components/Button/Button.stories.tsx

@@ -20,6 +20,11 @@ import { Meta, StoryFn } from "@storybook/react";
 import { Button as ButtonComponent } from "./Button";
 import VisibilityOnIcon from "@vector-im/compound-design-tokens/icons/visibility-on.svg";
 
+type Props = {
+  kind?: "primary" | "secondary" | "tertiary" | "destructive";
+  size?: "sm" | "lg";
+};
+
 export default {
   title: "Button",
   component: ButtonComponent,
@@ -29,13 +34,20 @@ export default {
       control: { type: "inline-radio" },
     },
   },
-  args: {},
+  args: {
+    size: "sm",
+  },
 } as Meta<typeof ButtonComponent>;
 
-const Template: StoryFn<typeof ButtonComponent> = (args) => (
+const Template: StoryFn<typeof ButtonComponent> = ({ kind, size }: Props) => (
   <div style={{ display: "flex", gap: 8 }}>
-    <ButtonComponent {...args}>Click me!</ButtonComponent>
-    <ButtonComponent Icon={VisibilityOnIcon} {...args}>
+    <ButtonComponent kind={kind} size={size}>
+      Click me!
+    </ButtonComponent>
+    <ButtonComponent Icon={VisibilityOnIcon} kind={kind} size={size}>
+      Click me!
+    </ButtonComponent>
+    <ButtonComponent as="a" href="#" kind={kind} size={size}>
       Click me!
     </ButtonComponent>
   </div>

src/components/Button/Button.tsx

@@ -15,15 +15,25 @@ limitations under the License.
 */
 
 import classNames from "classnames";
-import React, { ComponentType, PropsWithChildren } from "react";
+import React, {
+  ComponentType,
+  PropsWithChildren,
+  forwardRef,
+  ForwardedRef,
+  Ref,
+} from "react";
 import styles from "./Button.module.css";
 
-type ButtonProps<C extends React.ElementType> = {
-  /**
-   * The underlying HTML element to use. Can be a button or a link.
-   * @default "button"
-   */
-  as?: C;
+interface ButtonComponent {
+  // With the explicit `as` prop
+  <C extends React.ElementType>(
+    props: { as: C } & ButtonPropsFor<C>,
+  ): React.ReactElement;
+  // Without the explicit `as` prop, defaulting to a <button>
+  (props: ButtonPropsFor<"button">): React.ReactElement;
+}
+
+type ButtonOwnProps = PropsWithChildren<{
   /**
    * The type of button.
    */
@@ -32,30 +42,36 @@ type ButtonProps<C extends React.ElementType> = {
    * The t-shirt size of the button.
    */
   size?: "sm" | "lg";
-  /**
-   * The CSS class name.
-   */
-  className?: string;
   /**
    * An icon to display within the button.
    */
   Icon?: ComponentType<React.SVGAttributes<SVGElement>>;
-} & React.ComponentPropsWithoutRef<C>;
+}>;
+
+type ButtonPropsFor<C extends React.ElementType> = ButtonOwnProps &
+  Omit<React.ComponentPropsWithoutRef<C>, keyof ButtonOwnProps | "as"> & {
+    ref?: React.Ref<C>;
+  };
 
 /**
  * A button component that can be transformed into a link, but keep the button
  * styling using the `as` property.
  */
-export const Button = <C extends React.ElementType = "button">({
-  as,
-  kind = "primary",
-  size = "lg",
-  children,
-  className,
-  Icon,
-  ...props
-}: PropsWithChildren<ButtonProps<C>>): React.ReactElement => {
-  const Component = as || "button";
+export const Button = forwardRef(function Button<
+  C extends React.ElementType = "button",
+>(
+  {
+    as,
+    kind = "primary",
+    size = "lg",
+    children,
+    className,
+    Icon,
+    ...props
+  }: ButtonPropsFor<C> & { as?: C },
+  ref: ForwardedRef<C>,
+): React.ReactElement {
+  const Component = as || ("button" as const);
   const classes = classNames(styles.button, className, {
     [styles["has-icon"]]: Icon,
   });
@@ -64,6 +80,7 @@ export const Button = <C extends React.ElementType = "button">({
   return (
     <Component
       {...props}
+      ref={ref as Ref<C>}
       className={classes}
       data-kind={kind}
       data-size={size}
@@ -83,4 +100,4 @@ export const Button = <C extends React.ElementType = "button">({
       {children}
     </Component>
   );
-};
+}) as ButtonComponent;