From fd77ffb19e00fc85cc937fdb915463d73867ef8e Mon Sep 17 00:00:00 2001
From: Alexander Swed <lexswed@gmail.com>
Date: Sat, 31 Oct 2020 16:07:21 +0100
Subject: [PATCH] fix(Flex): extend Box for easier styling

---
 src/lib/Box/Box.tsx           |  1 -
 src/lib/Flex/Flex.tsx         |  5 +++--
 src/pages/components/Flex.mdx | 10 ++++++++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/src/lib/Box/Box.tsx b/src/lib/Box/Box.tsx
index da860d93..996ff079 100644
--- a/src/lib/Box/Box.tsx
+++ b/src/lib/Box/Box.tsx
@@ -27,7 +27,6 @@ const Box: React.FC<Props> = ({ children, css, ...props }) => {
       ),
     Object.values(props) //eslint-disable-line
   );
-
   return <Div css={Object.assign(style, css)} {...attrs} children={children} />;
 };
 
diff --git a/src/lib/Flex/Flex.tsx b/src/lib/Flex/Flex.tsx
index f8e2fa7b..1078c932 100644
--- a/src/lib/Flex/Flex.tsx
+++ b/src/lib/Flex/Flex.tsx
@@ -3,8 +3,9 @@ import React from 'react';
 import { StitchesProps, StitchesVariants } from '@stitches/react';
 import { styled } from '../stitches.config';
 import { gaps } from '../theme/variants';
+import Box from '../Box';
 
-const FlexBox = styled('div', {
+const FlexBox = styled(Box, {
   display: 'flex',
 
   variants: {
@@ -81,7 +82,7 @@ const FlexBox = styled('div', {
 
 const Flex = React.forwardRef<
   HTMLDivElement,
-  React.ComponentProps<typeof FlexBox> & { as?: keyof JSX.IntrinsicElements }
+  React.ComponentProps<typeof FlexBox> & { as?: keyof JSX.IntrinsicElements | React.ReactElement }
 >(({ space = 'none', flow = 'column', wrap = 'nowrap', main = 'start', cross = 'start', ...props }, ref) => (
   <FlexBox space={space} flow={flow} wrap={wrap} main={main} cross={cross} {...props} ref={ref} />
 ));
diff --git a/src/pages/components/Flex.mdx b/src/pages/components/Flex.mdx
index 63f8c0da..e07f9578 100644
--- a/src/pages/components/Flex.mdx
+++ b/src/pages/components/Flex.mdx
@@ -9,6 +9,16 @@ import { ExampleBox } from '../../pages-helpers';
 
 # Flex
 
+A `Box` with `flex` shorthands alignment properties.
+
+<Playground>
+  <Flex width="200px" bc="$orange050" cross="center" space="sm">
+    <ExampleBox />
+    <ExampleBox />
+    <ExampleBox />
+  </Flex>
+</Playground>
+
 ## Space
 
 Define gap between elements via `space` prop: