Inject Theme into Component

Button.js

// @flow

import React from 'react';
import classNames from 'classnames';

type ButtonTheme = {
    button: string,
    small: string,
    medium: string,
    large: string
};

type ButtonProps = {
    size: 'small' | 'medium' | 'large',
    children?: any,
    // Note that `theme` is required
    theme: ButtonTheme
};

const Button = (props: ButtonProps) => {
    const {
        size = 'large',
        children,
        theme
    } = props;

    const outputClassName = classNames(
        theme.button,
        theme[size]
    );
    return (
        <button className={outputClassName}>
            {children}
        </button>
    );
};

Button.displayName = 'Button';

export default Button;

mergeTheme.js

// @flow

import classnames from 'classnames';

type FunctionComponent<P, C> = (props: P, context: C) => ?React$Element<any>;
type ClassComponent<D, P, S> = Class<React$Component<D, P, S>>;
type ThemeType = { [className: string]: string };
type ThemePropType = { theme: ThemeType };

declare function mergeTheme<P: { theme: * }, V: $PropertyType<P, 'theme'>>(
    BaseComponent: ClassComponent<*, P, *>,
    injectedTheme: V
): FunctionComponent<$Diff<P, { theme: V }>, *>;

declare function mergeTheme<P: { theme: * }, V: $PropertyType<P, 'theme'>>(
    BaseComponent: FunctionComponent<P, *>,
    injectedTheme: V
): FunctionComponent<$Diff<P, { theme: V }>, *>;

function mergeTheme(BaseComponent, injectedTheme) {
    const ThemedComponent = ownProps => <BaseComponent {...ownProps} theme={injectedTheme} />;
    ThemedComponent.displayName = `Themed(${BaseComponent.displayName})`;
    return ThemedComponent;
}

export default mergeTheme;

PrimaryButton.js

// @flow

import Button from './Button';
import injectTheme from './injectTheme';
import theme from './styles/primary.scss';

// theme will be transpiled to look something like:
// {
//    button: primary-button--hx123,
//    small: primary-small--2j9r3,
//    medium: primary-medium--u9wjh,
//    large: primary-large--r92j2
// }

export default injectTheme(Button, theme);

test.js

import React from 'react';
import mergeTheme from './mergeTheme';

// == TEST CODE ==========================================

type SomeComponentTheme = {
    someClass: string
};

type SomeComponentProps = {
    someProp: string,
    theme: SomeComponentTheme
};

const SomeComponent = (props: SomeComponentProps) => {
    return <div className={props.theme.someClass}>{props.someProp}</div>;
};

const ThemedComponentOne = mergeTheme(SomeComponent, { someClass: 'poop' });
const ThemedComponentTwo = mergeTheme(SomeComponent, { anotherClass: 'poop '}); // should error but doesn't
const ThemedComponentThree = mergeTheme(SomeComponent, { someClass: 1 }); // should error but doesn't