From 6f9b502232161193be9c6556bdaf46cd18fd4225 Mon Sep 17 00:00:00 2001
From: James Brooks <jamesbrooks@spotify.com>
Date: Wed, 18 Jun 2025 17:49:05 +0100
Subject: [PATCH] Attempting to add forwardRef

Signed-off-by: James Brooks <jamesbrooks@spotify.com>
---
 .../canon/src/components/Button/Button.tsx    | 104 ++++++++++--------
 packages/canon/src/components/Button/types.ts |   9 +-
 packages/canon/src/types.ts                   |  20 ++--
 3 files changed, 75 insertions(+), 58 deletions(-)

diff --git a/packages/canon/src/components/Button/Button.tsx b/packages/canon/src/components/Button/Button.tsx
index 19e89171e7..905e81111b 100644
--- a/packages/canon/src/components/Button/Button.tsx
+++ b/packages/canon/src/components/Button/Button.tsx
@@ -15,54 +15,70 @@
  */
 
 import clsx from 'clsx';
+import {
+  ComponentPropsWithRef,
+  ElementType,
+  forwardRef,
+  ReactElement,
+  Ref,
+} from 'react';
+import { Button as RAButton } from 'react-aria-components';
 import { useResponsiveValue } from '../../hooks/useResponsiveValue';
 import type { ButtonProps } from './types';
 
 /** @public */
-export const Button = <C extends React.ElementType = 'button'>(
-  props: ButtonProps<C>,
-) => {
-  const {
-    as,
-    size = 'small',
-    variant = 'primary',
-    iconStart,
-    iconEnd,
-    children,
-    className,
-    ...rest
-  } = props;
+export const Button = forwardRef(
+  (props: ButtonProps<typeof RAButton>, ref: Ref<HTMLButtonElement>) => {
+    const {
+      as,
+      size = 'small',
+      variant = 'primary',
+      iconStart,
+      iconEnd,
+      children,
+      className,
+      ...rest
+    } = props;
 
-  const Component = as || 'button';
-  const responsiveSize = useResponsiveValue(size);
-  const responsiveVariant = useResponsiveValue(variant);
+    const Component = as || RAButton;
+    const responsiveSize = useResponsiveValue(size);
+    const responsiveVariant = useResponsiveValue(variant);
 
-  return (
-    <Component
-      className={clsx('canon-Button', className)}
-      data-variant={responsiveVariant}
-      data-size={responsiveSize}
-      {...rest}
-    >
-      {iconStart && (
-        <span
-          className="canon-ButtonIcon"
-          aria-hidden="true"
-          data-size={responsiveSize}
-        >
-          {iconStart}
-        </span>
-      )}
-      {children}
-      {iconEnd && (
-        <span
-          className="canon-ButtonIcon"
-          aria-hidden="true"
-          data-size={responsiveSize}
-        >
-          {iconEnd}
-        </span>
-      )}
-    </Component>
-  );
+    return (
+      <Component
+        className={clsx('canon-Button', className)}
+        data-variant={responsiveVariant}
+        data-size={responsiveSize}
+        ref={ref}
+        {...rest}
+      >
+        {iconStart && (
+          <span
+            className="canon-ButtonIcon"
+            aria-hidden="true"
+            data-size={responsiveSize}
+          >
+            {iconStart}
+          </span>
+        )}
+        {children}
+        {iconEnd && (
+          <span
+            className="canon-ButtonIcon"
+            aria-hidden="true"
+            data-size={responsiveSize}
+          >
+            {iconEnd}
+          </span>
+        )}
+      </Component>
+    );
+  },
+) as {
+  <TAs extends ElementType = typeof RAButton>(
+    props: ButtonProps<TAs> & { ref?: ComponentPropsWithRef<TAs>['ref'] },
+  ): ReactElement;
+  displayName: string;
 };
+
+Button.displayName = 'Button';
diff --git a/packages/canon/src/components/Button/types.ts b/packages/canon/src/components/Button/types.ts
index 1ff4e19cb0..327fa27bb6 100644
--- a/packages/canon/src/components/Button/types.ts
+++ b/packages/canon/src/components/Button/types.ts
@@ -15,17 +15,18 @@
  */
 
 import { Breakpoint } from '@backstage/canon';
-import { ReactElement } from 'react';
-import { PolymorphicComponentProp } from '../../types';
+import { ElementType, ReactElement, ReactNode } from 'react';
+import { PolymorphicComponentProps } from '../../types';
 
 /**
  * Properties for {@link Button}
  *
  * @public
  */
-export type ButtonProps<C extends React.ElementType> = PolymorphicComponentProp<
-  C,
+export type ButtonProps<TAs extends ElementType> = PolymorphicComponentProps<
+  TAs,
   {
+    children?: ReactNode;
     size?: 'small' | 'medium' | Partial<Record<Breakpoint, 'small' | 'medium'>>;
     variant?:
       | 'primary'
diff --git a/packages/canon/src/types.ts b/packages/canon/src/types.ts
index df3959e368..7e542f3635 100644
--- a/packages/canon/src/types.ts
+++ b/packages/canon/src/types.ts
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+import { ComponentPropsWithoutRef, ElementType } from 'react';
+
 /** @public */
 export type AsProps =
   | 'div'
@@ -133,19 +135,17 @@ export interface UtilityProps extends SpaceProps {
 }
 
 /** @public */
-export type AsProp<C extends React.ElementType> = {
-  as?: C;
+export type As<TAs extends ElementType> = {
+  as?: TAs;
 };
 
-/** @public */
-export type PropsToOmit<C extends React.ElementType, P> = keyof (AsProp<C> & P);
-
 /**
  * This is the first reusable type utility we built
  * @public
  */
-export type PolymorphicComponentProp<
-  C extends React.ElementType,
-  Props = {},
-> = React.PropsWithChildren<Props & AsProp<C>> &
-  Omit<React.ComponentPropsWithoutRef<C>, PropsToOmit<C, Props>>;
+export type PolymorphicComponentProps<
+  TAs extends ElementType,
+  TProps = {},
+> = TProps &
+  As<TAs> &
+  Omit<ComponentPropsWithoutRef<TAs>, keyof (As<TAs> & TProps)>;