refactor: Clean up unnecessary code

src/ContainerRender.js

@@ -1,6 +1,10 @@
 import React from 'react';
 import ReactDOM from 'react-dom';
 
+/**
+ * @deprecated Since we do not need support React15 any more.
+ * Will remove in next major version.
+ */
 export default class ContainerRender extends React.Component {
   static defaultProps = {
     autoMount: true,

src/Portal.tsx

@@ -1,10 +1,5 @@
 import * as React from 'react';
-import {
-  useRef,
-  useEffect,
-  forwardRef,
-  useImperativeHandle,
-} from 'react';
+import { useRef, useEffect, forwardRef, useImperativeHandle } from 'react';
 import ReactDOM from 'react-dom';
 import canUseDom from './Dom/canUseDom';
 
@@ -38,9 +33,9 @@ const Portal = forwardRef<PortalRef, PortalProps>((props, ref) => {
 
   useEffect(() => {
     return () => {
-      if (containerRef.current) {
-        containerRef.current.parentNode.removeChild(containerRef.current);
-      }
+      // [Legacy] This should not be handle by Portal but parent PortalWrapper instead.
+      // Since some component use `Portal` directly, we have to keep the logic here.
+      containerRef.current?.parentNode?.removeChild(containerRef.current);
     };
   }, []);
 

src/PortalWrapper.tsx

@@ -1,27 +1,25 @@
 /* eslint-disable no-underscore-dangle,react/require-default-props */
 import * as React from 'react';
-import * as ReactDOM from 'react-dom';
 import raf from './raf';
-import ContainerRender from './ContainerRender';
 import Portal, { PortalRef } from './Portal';
 import switchScrollingEffect from './switchScrollingEffect';
 import setStyle from './setStyle';
+import canUseDom from './Dom/canUseDom';
 
 let openCount = 0;
-const windowIsUndefined = !(
-  typeof window !== 'undefined' &&
-  window.document &&
-  window.document.createElement
-);
+const supportDom = canUseDom();
 
-const IS_REACT_16 = 'createPortal' in ReactDOM;
+/** @private Test usage only */
+export function getOpenCount() {
+  return process.env.NODE_ENV === 'test' ? openCount : 0;
+}
 
 // https://github.com/ant-design/ant-design/issues/19340
 // https://github.com/ant-design/ant-design/issues/19332
 let cacheOverflow = {};
 
 const getParent = (getContainer: GetContainer) => {
-  if (windowIsUndefined) {
+  if (!supportDom) {
     return null;
   }
   if (getContainer) {
@@ -66,7 +64,7 @@ class PortalWrapper extends React.Component<
 > {
   container?: HTMLElement;
 
-  _component?: PortalRef;
+  componentRef: React.RefObject<PortalRef> = React.createRef();
 
   rafId?: number;
 
@@ -76,12 +74,10 @@ class PortalWrapper extends React.Component<
     visible: boolean;
   }) => void;
 
-  removeContainer?: Function;
-
   constructor(props: PortalWrapperProps) {
     super(props);
     const { visible, getContainer } = props;
-    if (!windowIsUndefined && getParent(getContainer) === document.body) {
+    if (supportDom && getParent(getContainer) === document.body) {
       openCount = visible ? openCount + 1 : openCount;
     }
     this.state = {
@@ -104,11 +100,11 @@ class PortalWrapper extends React.Component<
 
   componentWillUnmount() {
     const { visible, getContainer } = this.props;
-    if (!windowIsUndefined && getParent(getContainer) === document.body) {
+    if (supportDom && getParent(getContainer) === document.body) {
       // 离开时不会 render， 导到离开时数值不变，改用 func 。。
       openCount = visible && openCount ? openCount - 1 : openCount;
     }
-    this.removeCurrentContainer(visible);
+    this.removeCurrentContainer();
     raf.cancel(this.rafId);
   }
 
@@ -121,7 +117,7 @@ class PortalWrapper extends React.Component<
       } = prevProps;
       if (
         visible !== prevVisible &&
-        !windowIsUndefined &&
+        supportDom &&
         getParent(getContainer) === document.body
       ) {
         openCount = visible && !prevVisible ? openCount + 1 : openCount - 1;
@@ -134,7 +130,7 @@ class PortalWrapper extends React.Component<
           ? getContainer.toString() !== prevGetContainer.toString()
           : getContainer !== prevGetContainer
       ) {
-        _self.removeCurrentContainer(false);
+        _self.removeCurrentContainer();
       }
     }
     return {
@@ -157,7 +153,7 @@ class PortalWrapper extends React.Component<
   };
 
   getContainer = () => {
-    if (windowIsUndefined) {
+    if (!supportDom) {
       return null;
     }
     if (!this.container) {
@@ -179,26 +175,10 @@ class PortalWrapper extends React.Component<
     }
   };
 
-  savePortal = (c: PortalRef) => {
-    // Warning: don't rename _component
-    // https://github.com/react-component/util/pull/65#discussion_r352407916
-    this._component = c;
-  };
-
-  removeCurrentContainer = visible => {
-    this.container = null;
-    this._component = null;
-    if (!IS_REACT_16) {
-      if (visible) {
-        this.renderComponent({
-          afterClose: this.removeContainer,
-          onClose() {},
-          visible: false,
-        });
-      } else {
-        this.removeContainer();
-      }
-    }
+  removeCurrentContainer = () => {
+    // Portal will remove from `parentNode`.
+    // Let's handle this again to avoid refactor issue.
+    this.container?.parentNode?.removeChild(this.container);
   };
 
   /**
@@ -233,34 +213,10 @@ class PortalWrapper extends React.Component<
       getContainer: this.getContainer,
       switchScrollingEffect: this.switchScrollingEffect,
     };
-    // support react15
-    if (!IS_REACT_16) {
-      return (
-        <ContainerRender
-          parent={this}
-          visible={visible}
-          autoDestroy={false}
-          getComponent={(extra = {}) =>
-            children({
-              ...extra,
-              ...childProps,
-              ref: this.savePortal,
-            })
-          }
-          getContainer={this.getContainer}
-          forceRender={forceRender}
-        >
-          {({ renderComponent, removeContainer }) => {
-            this.renderComponent = renderComponent;
-            this.removeContainer = removeContainer;
-            return null;
-          }}
-        </ContainerRender>
-      );
-    }
-    if (forceRender || visible || this._component) {
+
+    if (forceRender || visible || this.componentRef.current) {
       portal = (
-        <Portal getContainer={this.getContainer} ref={this.savePortal}>
+        <Portal getContainer={this.getContainer} ref={this.componentRef}>
           {children(childProps)}
         </Portal>
       );

tests/Portal.test.tsx

@@ -1,11 +1,16 @@
 import React from 'react';
 import { mount } from 'enzyme';
-import PortalWrapper from '../src/PortalWrapper';
+import { act } from 'react-dom/test-utils';
+import PortalWrapper, { getOpenCount } from '../src/PortalWrapper';
 import Portal from '../src/Portal';
 
 describe('Portal', () => {
   let container: HTMLDivElement;
 
+  // Mock for raf
+  window.requestAnimationFrame = callback => window.setTimeout(callback);
+  window.cancelAnimationFrame = id => window.clearTimeout(id);
+
   beforeEach(() => {
     container = document.createElement('div');
     document.body.appendChild(container);
@@ -46,4 +51,107 @@ describe('Portal', () => {
     wrapper.setProps({ justForceUpdate: true });
     expect(didUpdate).toHaveBeenCalledTimes(2);
   });
+
+  describe('getContainer', () => {
+    it('string', () => {
+      const div = document.createElement('div');
+      div.id = 'bamboo-light';
+      document.body.appendChild(div);
+
+      mount(
+        <PortalWrapper visible getContainer="#bamboo-light">
+          {() => <div>2333</div>}
+        </PortalWrapper>,
+      );
+
+      expect(document.querySelector('#bamboo-light').childElementCount).toEqual(
+        1,
+      );
+
+      document.body.removeChild(div);
+    });
+
+    it('function', () => {
+      const div = document.createElement('div');
+
+      mount(
+        <PortalWrapper visible getContainer={() => div}>
+          {() => <div>2333</div>}
+        </PortalWrapper>,
+      );
+
+      expect(div.childElementCount).toEqual(1);
+    });
+
+    it('htmlElement', () => {
+      const div = document.createElement('div');
+
+      mount(
+        <PortalWrapper visible getContainer={div}>
+          {() => <div>2333</div>}
+        </PortalWrapper>,
+      );
+
+      expect(div.childElementCount).toEqual(1);
+    });
+
+    it('delay', () => {
+      jest.useFakeTimers();
+      const divRef = React.createRef<HTMLDivElement>();
+      const wrapper = mount(
+        <div>
+          <PortalWrapper visible getContainer={() => divRef.current}>
+            {() => <div />}
+          </PortalWrapper>
+          <div ref={divRef} />
+        </div>,
+      );
+
+      act(() => {
+        jest.runAllTimers();
+        wrapper.update();
+      });
+
+      expect(divRef.current.childElementCount).toEqual(1);
+      jest.useRealTimers();
+    });
+  });
+
+  it('openCount', () => {
+    const Demo = ({ count, visible }: { count: number; visible: boolean }) => {
+      return (
+        <>
+          {new Array(count).fill(null).map((_, index) => (
+            <PortalWrapper key={index} visible={visible}>
+              {() => <div>2333</div>}
+            </PortalWrapper>
+          ))}
+        </>
+      );
+    };
+
+    const wrapper = mount(<Demo count={1} visible />);
+    expect(getOpenCount()).toEqual(1);
+
+    wrapper.setProps({ count: 2 });
+    expect(getOpenCount()).toEqual(2);
+
+    wrapper.setProps({ count: 1 });
+    expect(getOpenCount()).toEqual(1);
+
+    wrapper.setProps({ visible: false });
+    expect(getOpenCount()).toEqual(0);
+  });
+
+  it('wrapperClassName', () => {
+    const wrapper = mount(
+      <PortalWrapper visible wrapperClassName="bamboo">
+        {() => <div />}
+      </PortalWrapper>,
+    );
+    expect((wrapper.instance() as any).container.className).toEqual('bamboo');
+
+    wrapper.setProps({ wrapperClassName: 'light' });
+    expect((wrapper.instance() as any).container.className).toEqual('light');
+  });
 });