withErrorBoundary forwards refs

Author: bvaughn

src/withErrorBoundary.test.tsx

@@ -2,7 +2,7 @@
  * @jest-environment jsdom
  */
 
-import { PropsWithChildren } from "react";
+import { Component, createRef, PropsWithChildren } from "react";
 import { createRoot } from "react-dom/client";
 import { act } from "react-dom/test-utils";
 import { withErrorBoundary } from "./withErrorBoundary";
@@ -53,4 +53,28 @@ describe("withErrorBoundary", () => {
     render();
     expect(container.textContent).toBe("Error");
   });
+
+  it("should forward refs", () => {
+    type Props = { foo: string };
+
+    class Inner extends Component<Props> {
+      test() {}
+      render() {
+        return this.props.foo;
+      }
+    }
+
+    const Wrapped = withErrorBoundary(Inner, {
+      fallback: <div>Error</div>,
+    });
+
+    const ref = createRef<Inner>();
+
+    act(() => {
+      root.render(<Wrapped foo="abc" ref={ref} />);
+    });
+
+    expect(ref.current).not.toBeNull();
+    expect(typeof ref.current?.test).toBe("function");
+  });
 });

src/withErrorBoundary.ts

@@ -1,21 +1,30 @@
-import { ComponentType, createElement } from "react";
+import {
+  createElement,
+  forwardRef,
+  ForwardedRef,
+  RefAttributes,
+  ForwardRefExoticComponent,
+  PropsWithoutRef,
+  ComponentType,
+} from "react";
 import { ErrorBoundary } from "./ErrorBoundary";
 import { ErrorBoundaryProps } from "./types";
 
 export function withErrorBoundary<Props extends Object>(
-  Component: ComponentType<Props>,
+  component: ComponentType<Props>,
   errorBoundaryProps: ErrorBoundaryProps
-): ComponentType<Props> {
-  const Wrapped: ComponentType<Props> = (props: Props) => {
-    return createElement(
-      ErrorBoundary,
-      errorBoundaryProps,
-      createElement(Component, props)
-    );
-  };
+): ForwardRefExoticComponent<PropsWithoutRef<Props> & RefAttributes<any>> {
+  const Wrapped = forwardRef<ComponentType<Props>, Props>(
+    (props: Props, ref: ForwardedRef<ComponentType<Props>>) =>
+      createElement(
+        ErrorBoundary,
+        errorBoundaryProps,
+        createElement(component, { ...props, ref })
+      )
+  );
 
   // Format for display in DevTools
-  const name = Component.displayName || Component.name || "Unknown";
+  const name = component.displayName || component.name || "Unknown";
   Wrapped.displayName = `withErrorBoundary(${name})`;
 
   return Wrapped;