From 57baa16e601e1f7152403f5e4c592632e6f9cfcf Mon Sep 17 00:00:00 2001
From: Jon Lambert <jonlambert@icloud.com>
Date: Wed, 27 Apr 2022 11:20:10 +0100
Subject: [PATCH] feat: support passing refs to intrinsic elements

---
 src/index.ts         | 20 +++++++++++++++++++-
 src/propose.test.tsx | 12 ++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index a865591..aa392d6 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -2,6 +2,10 @@ import React, {
   createElement,
   JSXElementConstructor,
   ComponentProps,
+  forwardRef,
+  ForwardedRef,
+  RefCallback,
+  Ref,
 } from 'react';
 
 type IntrinsicElementKeys = keyof JSX.IntrinsicElements;
@@ -9,11 +13,25 @@ type IntrinsicElementKeys = keyof JSX.IntrinsicElements;
 export function propose<
   ComponentType extends JSXElementConstructor<any> | IntrinsicElementKeys,
   OriginalProps extends ComponentProps<ComponentType>,
-  SuppliedProps extends Partial<OriginalProps>
+  SuppliedProps extends Partial<OriginalProps>,
+  RefType extends ComponentType extends IntrinsicElementKeys
+    ?
+        | Ref<JSX.IntrinsicElements[ComponentType]>
+        | RefCallback<JSX.IntrinsicElements[ComponentType]>
+    : "Unable to pass a ref to a function component"
 >(Component: ComponentType, props: SuppliedProps, displayName?: string) {
   type FinalProps = Omit<OriginalProps, keyof SuppliedProps> &
     Partial<SuppliedProps>;
 
+  if (typeof Component === 'string') {
+    const NewComponent = forwardRef<RefType, FinalProps>((p, ref) => {
+      const combinedProps = { ...props, ...p, ref };
+      return createElement(Component, combinedProps);
+    });
+    NewComponent.displayName = displayName;
+    return NewComponent;
+  }
+
   const NewComponent: React.FC<FinalProps> = (p) => {
     const combinedProps = { ...props, ...p };
     return createElement(Component, combinedProps);
diff --git a/src/propose.test.tsx b/src/propose.test.tsx
index da89181..e199a10 100644
--- a/src/propose.test.tsx
+++ b/src/propose.test.tsx
@@ -1,8 +1,8 @@
 import 'jest';
 import '@testing-library/jest-dom/extend-expect';
 
-import React from 'react';
-import { render } from '@testing-library/react';
+import React, { createRef, useRef } from 'react';
+import { act, render } from '@testing-library/react';
 import { propose } from '.';
 
 describe('propose', () => {
@@ -28,6 +28,7 @@ describe('propose', () => {
 
     const Decorated = propose(Base, { title });
     const { getByRole, getByText } = render(<Decorated message={message} />);
+
     expect(getByRole('heading')).toHaveTextContent(title);
     expect(getByText(message)).toBeInTheDocument();
   });
@@ -73,4 +74,11 @@ describe('propose', () => {
     const Derived = propose(Base, { variant: 'primary' });
     render(<Derived onClick={() => {}} variant="primary" />);
   });
+
+  it('accepts and passes down ref', () => {
+    const Decorated = propose('button', {});
+    const ref = createRef<HTMLButtonElement>();
+    render(<Decorated ref={ref as any}>Hello, world</Decorated>);
+    expect(ref.current?.innerHTML).toEqual('Hello, world');
+  });
 });