chore: support ref (#5)

Author: zombieJ

package.json

@@ -39,7 +39,7 @@
     "test": "npm run test:client && npm run test:server",
     "test:client": "umi-test --testPathIgnorePatterns=ssr.test.tsx --testPathIgnorePatterns=ssr.test.tsx",
     "test:coverage": "npm run test:client --coverage",
-    "test:server": "umi-test --env=node tests/ssr.test.tsx",
+    "test:server": "umi-test --env=node -- tests/ssr.test.tsx",
     "watch": "father dev"
   },
   "dependencies": {

src/Portal.tsx

@@ -1,6 +1,7 @@
 import * as React from 'react';
 import { createPortal } from 'react-dom';
 import canUseDom from 'rc-util/lib/Dom/canUseDom';
+import { supportRef, useComposeRef } from 'rc-util/lib/ref';
 import OrderContext from './Context';
 import useDom from './useDom';
 import useScrollLocker from './useScrollLocker';
@@ -47,7 +48,7 @@ const getPortalContainer = (getContainer: GetContainer) => {
   return getContainer;
 };
 
-export default function Portal(props: PortalProps) {
+const Portal = React.forwardRef<any, PortalProps>((props, ref) => {
   const {
     open,
     autoLock,
@@ -93,6 +94,15 @@ export default function Portal(props: PortalProps) {
         mergedContainer === document.body),
   );
 
+  // =========================== Ref ===========================
+  let childRef: React.Ref<any> = null;
+
+  if (children && supportRef(children) && ref) {
+    ({ ref: childRef } = children as any);
+  }
+
+  const mergedRef = useComposeRef(childRef, ref);
+
   // ========================= Render ==========================
   // Do not render when nothing need render
   // When innerContainer is `undefined`, it may not ready since user use ref in the same render
@@ -103,9 +113,24 @@ export default function Portal(props: PortalProps) {
   // Render inline
   const renderInline = mergedContainer === false || inlineMock();
 
+  let reffedChildren = children;
+  if (ref) {
+    reffedChildren = React.cloneElement(children as any, {
+      ref: mergedRef,
+    });
+  }
+
   return (
     <OrderContext.Provider value={queueCreate}>
-      {renderInline ? children : createPortal(children, mergedContainer)}
+      {renderInline
+        ? reffedChildren
+        : createPortal(reffedChildren, mergedContainer)}
     </OrderContext.Provider>
   );
+});
+
+if (process.env.NODE_ENV !== 'production') {
+  Portal.displayName = 'Portal';
 }
+
+export default Portal;

tests/index.test.tsx

@@ -172,4 +172,52 @@ describe('Portal', () => {
     rerender(renderDemo(false, true));
     expect(document.querySelector('p')).toBeFalsy();
   });
+
+  describe.only('ref-able', () => {
+    it('support forwardRef', () => {
+      const elementRef = React.createRef<HTMLParagraphElement>();
+      const portalRef = React.createRef();
+
+      render(
+        <Portal ref={portalRef} open>
+          <p ref={elementRef}>Bamboo</p>
+        </Portal>,
+      );
+
+      expect(elementRef.current).toBe(document.querySelector('p'));
+      expect(portalRef.current).toBe(document.querySelector('p'));
+    });
+
+    it('not support fragment', () => {
+      const elementRef = React.createRef<HTMLParagraphElement>();
+      const portalRef = React.createRef();
+
+      render(
+        <Portal ref={portalRef} open>
+          <>
+            <p ref={elementRef}>Bamboo</p>
+          </>
+        </Portal>,
+      );
+
+      expect(elementRef.current).toBe(document.querySelector('p'));
+      expect(portalRef.current).toBeFalsy();
+    });
+
+    it('not support FC', () => {
+      const elementRef = React.createRef<HTMLParagraphElement>();
+      const portalRef = React.createRef();
+
+      const P = (props: any) => <p {...props} />;
+
+      render(
+        <Portal ref={portalRef} open>
+          <P ref={elementRef}>Bamboo</P>
+        </Portal>,
+      );
+
+      expect(elementRef.current).toBeFalsy();
+      expect(portalRef.current).toBeFalsy();
+    });
+  });
 });