feat: add generic type for second theme

src/index.ts

@@ -7,12 +7,14 @@ type ThemeFn<T> = ThemeFn1<T> | ThemeFn2<T>
 type ThemeFn1<T> = (primitives: T) => ThemeTokens
 type ThemeFn2<T> = (primitives: T) => ThemeFn1<T>
 
-// TODO: theme args must be fn or plain shape.
-export function withTokens<T extends Shape<any>>(themeFn1: ThemeFn<T>, themeFn2?: ThemeFn<T>) {
-  return (primitives1?: T) => (primitives2?: T) => {
-    const composedPrimitives = deepmerge<T>(primitives1 || {}, primitives2 || {})
-    const theme1 = esModuleInterop<ThemeFn1<T>>(themeFn1 as any)
-    const theme2 = esModuleInterop<ThemeFn1<T>>(themeFn2 as any)
+export function withTokens<T extends Shape<any>, U extends Shape<any> = Shape<any>>(
+  themeFn1: ThemeFn<T>,
+  themeFn2?: ThemeFn<U>,
+) {
+  return (primitives1?: T) => (primitives2?: U) => {
+    const composedPrimitives = deepmerge(primitives1 || {}, primitives2 || {})
+    const theme1 = esModuleInterop<ThemeFn1<any>>(themeFn1 as any)
+    const theme2 = esModuleInterop<ThemeFn1<any>>(themeFn2 as any)
     if (theme2 === undefined) {
       return theme1(composedPrimitives)
     }