createFocusWrapper.js

import React from "react";

function compareNodePosition(a, b) {
  if (!a || !b) return "none";

  const documentPosition = a.compareDocumentPosition(b);

  if (documentPosition & Node.DOCUMENT_POSITION_CONTAINS) return "contains";
  if (documentPosition & Node.DOCUMENT_POSITION_CONTAINED_BY)
    return "contained_by";
  if (documentPosition & Node.DOCUMENT_POSITION_PRECEDING) return "preceding";
  if (documentPosition & Node.DOCUMENT_POSITION_FOLLOWING) return "following";

  return "none";
}

export default function createFocusWrapper(WrappedComponent) {
  return class FocusWrapper extends React.Component {
    state = {
      tabIndex: 0
    };

    componentDidMount() {
      document.addEventListener("focus", this._handleDocumentFocus, true);
    }

    componentWillUnmount() {
      document.removeEventListener("focus", this._handleDocumentFocus, true);
    }

    _previousActiveElement = null;

    _backupRef = React.createRef();

    _getRef() {
      const { forwardedRef } = this.props;

      return forwardedRef || this._backupRef;
    }

    _handleDocumentFocus = e => {
      const ref = this._getRef();

      if (ref.current) {
        this.setState({
          tabIndex: ref.current.contains(e.target) ? -1 : 0
        });
      }

      this._previousActiveElement = e.target;
    };

    _handleFocus = e => {
      const { onFocusFirst, onFocusLast, onFocus } = this.props;

      const ref = this._getRef();

      if (e.target === ref.current) {
        const position = compareNodePosition(
          document.activeElement,
          this._previousActiveElement
        );

        switch (position) {
          case "preceding":
          case "none":
            onFocusFirst && onFocusFirst();

            e.stopPropagation();
            e.preventDefault();
            break;
          case "following":
            onFocusLast && onFocusLast();

            e.stopPropagation();
            e.preventDefault();
            break;
          default:
            break;
        }
      } else {
        onFocus && onFocus();
      }
    };

    render() {
      const { forwardedRef, onFocusFirst, onFocusLast, ...rest } = this.props;
      const { tabIndex } = this.state;

      return (
        <WrappedComponent
          {...rest}
          tabIndex={tabIndex}
          ref={this._getRef()}
          onFocus={this._handleFocus}
        />
      );
    }
  };
}