3 files changed, 11 insertions, 37 deletions
diff --git a/packages/instant/src/components/ui/container.tsx b/packages/instant/src/components/ui/container.tsx
index 36d7cf2ec..0e518be88 100644
--- a/packages/instant/src/components/ui/container.tsx
+++ b/packages/instant/src/components/ui/container.tsx
@@ -1,7 +1,7 @@
 import * as _ from 'lodash';
 import { darken } from 'polished';
 
-import { media } from '../../style/media';
+import { MediaChoice, stylesForMedia } from '../../style/media';
 import { ColorOption, styled } from '../../style/theme';
 import { cssRuleIfExists } from '../../style/util';
 
@@ -12,8 +12,8 @@ export interface ContainerProps {
     right?: string;
     bottom?: string;
     left?: string;
-    width?: string;
-    height?: string;
+    width?: MediaChoice;
+    height?: MediaChoice;
     maxWidth?: string;
     margin?: string;
     marginTop?: string;
@@ -36,22 +36,8 @@ export interface ContainerProps {
     overflow?: string;
     darkenOnHover?: boolean;
     flexGrow?: string | number;
-
-    smallWidth?: string;
-    smallHeight?: string;
 }
 
-const mediaStyles = (props: ContainerProps) => {
-    if (!_.some([props.smallWidth, props.smallHeight])) {
-        return '';
-    }
-
-    return media.small`
-        width: ${props.smallWidth || props.width || 'auto'}
-        height: ${props.smallHeight || props.height || 'auto'}
-    `;
-};
-
 // TODO Dont commit flex grow
 export const Container =
     styled.div <
@@ -84,7 +70,8 @@ export const Container =
     ${props => cssRuleIfExists(props, 'cursor')}
     ${props => cssRuleIfExists(props, 'overflow')}
     ${props => (props.hasBoxShadow ? `box-shadow: 0px 2px 10px rgba(0, 0, 0, 0.1)` : '')};
-    ${props => mediaStyles(props)}
+    ${props => stylesForMedia('width', props.width || 'auto')}
+    ${props => stylesForMedia('height', props.height || 'auto')}
     background-color: ${props => (props.backgroundColor ? props.theme[props.backgroundColor] : 'none')};
     border-color: ${props => (props.borderColor ? props.theme[props.borderColor] : 'none')};
     &:hover {
diff --git a/packages/instant/src/components/ui/flex.tsx b/packages/instant/src/components/ui/flex.tsx
index f85d45e36..c3d643a79 100644
--- a/packages/instant/src/components/ui/flex.tsx
+++ b/packages/instant/src/components/ui/flex.tsx
@@ -1,6 +1,6 @@
 import * as _ from 'lodash';
 
-import { media } from '../../style/media';
+import { MediaChoice, stylesForMedia } from '../../style/media';
 import { ColorOption, styled } from '../../style/theme';
 import { cssRuleIfExists } from '../../style/util';
 
@@ -9,27 +9,13 @@ export interface FlexProps {
     flexWrap?: 'wrap' | 'nowrap';
     justify?: 'flex-start' | 'center' | 'space-around' | 'space-between' | 'space-evenly' | 'flex-end';
     align?: 'flex-start' | 'center' | 'space-around' | 'space-between' | 'space-evenly' | 'flex-end';
-    width?: string;
-    height?: string;
+    width?: MediaChoice;
+    height?: MediaChoice;
     backgroundColor?: ColorOption;
     inline?: boolean;
     flexGrow?: number | string;
-
-    smallWidth?: string;
-    smallHeight?: string;
 }
 
-const mediaStyles = (props: FlexProps) => {
-    if (!_.some([props.smallWidth, props.smallHeight])) {
-        return '';
-    }
-
-    return media.small`
-        width: ${props.smallWidth || props.width || 'auto'}
-        height: ${props.smallHeight || props.height || 'auto'}
-    `;
-};
-
 export const Flex =
     styled.div <
     FlexProps >
@@ -43,7 +29,8 @@ export const Flex =
     ${props => cssRuleIfExists(props, 'width')}
     ${props => cssRuleIfExists(props, 'height')}
     background-color: ${props => (props.backgroundColor ? props.theme[props.backgroundColor] : 'none')};
-    ${props => mediaStyles(props)}
+    ${props => stylesForMedia('width', props.width || 'auto')}
+    ${props => stylesForMedia('height', props.height || 'auto')}
 `;
 
 Flex.defaultProps = {
diff --git a/packages/instant/src/components/ui/overlay.tsx b/packages/instant/src/components/ui/overlay.tsx
index 12bde1f2b..0fde995e0 100644
--- a/packages/instant/src/components/ui/overlay.tsx
+++ b/packages/instant/src/components/ui/overlay.tsx
@@ -18,7 +18,7 @@ const PlainOverlay: React.StatelessComponent<OverlayProps> = ({ children, classN
         <Container position="absolute" top="0px" right="0px">
             <Icon height={18} width={18} color={ColorOption.white} icon="closeX" onClick={onClose} padding="2em 2em" />
         </Container>
-        <Container smallWidth="100%" smallHeight="100%">
+        <Container width={{ default: 'auto', sm: '100%' }} height={{ default: 'auto', sm: '100%' }}>
             {children}
         </Container>
     </Flex>