React async tree walker

walk.js

/* @flow */
import * as React from 'react';

type ReactContext = Object;

/*
 * Walk a react element.
 */
async function walkTree(
    element: React.Element<*>,
    visitor: (
        el: React.Element<*>,
        instance: ?React.Component<*, *>,
        context: ReactContext
    ) => boolean | Promise<boolean>,
    context: ReactContext
): Promise<void> {
    if (Array.isArray(element)) {
        await Promise.all(
            element.map(child => walkTree(child, visitor, context))
        );

        return;
    }

    if (!element) {
        return;
    }

    const Component = element.type;

    // a stateless functional component or a class
    if (typeof Component === 'function') {
        const props = Object.assign({}, Component.defaultProps, element.props);
        let child;
        let childContext = context;

        // Are we are a react class?
        if (Component.prototype && Component.prototype.isReactComponent) {
            const instance = new Component(props, context);

            // In case the user doesn't pass these to super in the constructor
            instance.props = instance.props || props;
            instance.context = instance.context || context;
            // set the instance state to null (not undefined) if not set, to match React behaviour
            instance.state = instance.state || null;

            // Override setState to just change the state, not queue up an update.
            //   (we can't do the default React thing as we aren't mounted "properly"
            //   however, we don't need to re-render as well only support setState in
            //   componentWillMount, which happens *before* render).
            instance.setState = newStateFn => {
                let newState = newStateFn;

                if (typeof newStateFn === 'function') {
                    newState = newStateFn(
                        instance.state,
                        instance.props,
                        instance.context
                    );
                }
                instance.state = Object.assign({}, instance.state, newState);
            };

            if (instance.componentWillMount) {
                instance.componentWillMount();
            }

            if (instance.getChildContext) {
                childContext = Object.assign(
                    {},
                    context,
                    instance.getChildContext()
                );
            }

            if ((await visitor(element, instance, context)) === false) {
                return;
            }

            child = instance.render();
        } else {
            // just a stateless functional
            if ((await visitor(element, null, context)) === false) {
                return;
            }

            child = Component(props, context);
        }

        // Traverse children
        await walkTree(child, visitor, childContext);
    } else {
        // a basic string or dom element, just get children
        if ((await visitor(element, null, context)) === false) {
            return;
        }

        if (element.props && element.props.children) {
            const pending = [];

            React.Children.forEach(
                element.props.children,
                (child: React.Element<*>) => {
                    pending.push(walkTree(child, visitor, context));
                }
            );

            await Promise.all(pending);
        }
    }
}

export default walkTree;