diff --git a/src/index.ts b/src/index.ts
index cdd7ed8..8dbda55 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -4,9 +4,9 @@ import { Color } from 'colorette';
 
 export { colorette as color };
 
-export type ColorTheme = Record<TokenType, Color>;
+export type ColorTheme = Partial<Record<TokenType, Color | ((token: string) => string)>>;
 
-const defaultTheme: ColorTheme = {
+const defaultTheme: Record<TokenType, Color> = {
   Whitespace: colorette.gray,
   Brace: colorette.gray,
   Bracket: colorette.gray,