1 files changed, 23 insertions, 15 deletions
diff --git a/packages/instant/src/style/media.ts b/packages/instant/src/style/media.ts
index beabbac46..5e7aaba37 100644
--- a/packages/instant/src/style/media.ts
+++ b/packages/instant/src/style/media.ts
@@ -14,30 +14,38 @@ const generateMediaWrapper = (screenWidth: ScreenWidths) => (...args: any[]) =>
     }
 `;
 
-const media = {
+export const media = {
     small: generateMediaWrapper(ScreenWidths.Sm),
     medium: generateMediaWrapper(ScreenWidths.Md),
     large: generateMediaWrapper(ScreenWidths.Lg),
 };
 
-export interface ScreenSpecifications {
-    default: string;
-    sm?: string;
-    md?: string;
-    lg?: string;
+export interface ScreenSpecification<T> {
+    default: T;
+    sm?: T;
+    md?: T;
+    lg?: T;
 }
-export type MediaChoice = string | ScreenSpecifications;
-export const stylesForMedia = (cssPropertyName: string, choice: MediaChoice): InterpolationValue[] => {
-    if (typeof choice === 'string') {
+export type OptionallyScreenSpecific<T> = T | ScreenSpecification<T>;
+export type MediaChoice = OptionallyScreenSpecific<string>;
+/**
+ * Given a css property name and a OptionallyScreenSpecific value,
+ * generates css properties with screen-specific viewport styling
+ */
+export function stylesForMedia<T extends string | number>(
+    cssPropertyName: string,
+    choice: OptionallyScreenSpecific<T>,
+): InterpolationValue[] {
+    if (typeof choice === 'object') {
         return css`
-            ${cssPropertyName}: ${choice};
-        `;
-    }
-
-    return css`
         ${cssPropertyName}: ${choice.default};
         ${choice.lg && media.large`${cssPropertyName}: ${choice.lg}`}
         ${choice.md && media.medium`${cssPropertyName}: ${choice.md}`}
         ${choice.sm && media.small`${cssPropertyName}: ${choice.sm}`}
     `;
-};
+    } else {
+        return css`
+            ${cssPropertyName}: ${choice};
+        `;
+    }
+}