React Flex Component with Typescript

I love wrapper instead of long inline style jsx syntax. This is totally useful for me and save many line of codes, so instead of

<div style={{ display: "flex", flexDirection: "row", /*...*/ }}>
    {/*...*/}
</div>

I made a wrapper for it (this is called Polymorphic Component, I guess?)

import React, { ElementType, ForwardedRef, forwardRef, MutableRefObject, ReactElement } from "react";

type FlexProps<T extends React.ElementType> = React.ComponentPropsWithoutRef<T> & {
    as?: T;
    children?: React.ReactNode;
    style?: React.CSSProperties;
    ref?: MutableRefObject<HTMLElement> | ForwardedRef<HTMLElement>;
};

type PolymorphicComponent = <T extends ElementType = "div">(props: FlexProps<T>) => ReactElement | null;
type PolymorphicComponentWithDisplayName = PolymorphicComponent & { displayName?: string };

export const Flex: PolymorphicComponentWithDisplayName = forwardRef(
    <T extends ElementType>(props: FlexProps<T>, ref) => {
        const { as, children, style, ...rest } = props;
        const Element = as || "div";
        return (
            <Element ref={ref} style={{ display: "flex", ...style }} {...rest}>
                {children}
            </Element>
        );
    }
);
Flex.displayName = "Flex";

What makes it even cooler is, you can pass as to render the wrapper component as another component without losing type checking. So if you want to render a span instead of div you can do so by passing as={"span"}, or if you want it as a motion.div then you can as={motion.div} and still get all motion props in the implementation.

Now, to make it even simpler, we can also wrap the mostly used flex style in a component, for example

export const FlexRowAlignCenter: PolymorphicComponentWithDisplayName = forwardRef(
    <T extends ElementType>(props: FlexProps<T>, ref) => {
        const { style, children, ...rest } = props;
        return (
            <Flex ref={ref} style={{ display: "flex", flexDirection: "row", alignItems: "center", ...style }} {...rest}>
                {children}
            </Flex>
        );
    }
);
FlexRowAlignCenter.displayName = "FlexRowAlignCenter";

export const FlexColumn: PolymorphicComponentWithDisplayName = forwardRef(
    <T extends ElementType>(props: FlexProps<T>, ref) => {
        const { style, children, ...rest } = props;
        return (
            <Flex ref={ref} style={{ display: "flex", flexDirection: "column", ...style }} {...rest}>
                {children}
            </Flex>
        );
    }
);
FlexColumn.displayName = "FlexColumn";

Both are mostly used flex container in my usage. To apply it then simply

<FlexRowAlignCenter style={{ gap: "2rem" }}>
    <button>a</button>
    <button>b</button>
    <button>c</button>
    <button>d</button>
    <button>e</button>
</FlexRowAlignCenter>

That's it, have fun 🐳🌊🐠