From 067f595f0c7581d2eed247c19cae602885dbde23 Mon Sep 17 00:00:00 2001
From: zombiej <smith3816@gmail.com>
Date: Mon, 21 Mar 2022 14:48:55 +0800
Subject: [PATCH] feat: Add formatToken support

---
 src/useCacheToken.tsx | 24 ++++++++++++++++++++----
 tests/index.spec.tsx  | 14 ++++++++++++--
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/src/useCacheToken.tsx b/src/useCacheToken.tsx
index 59a8d902..e18552c4 100644
--- a/src/useCacheToken.tsx
+++ b/src/useCacheToken.tsx
@@ -7,13 +7,24 @@ import { flattenToken, token2key } from './util';
 
 const EMPTY_OVERRIDE = {};
 
-export interface Option {
+export interface Option<DerivativeToken> {
   /**
    * Generate token with salt.
    * This is used to generate different hashId even same derivative token for different version.
    */
   salt?: string;
   override?: object;
+  /**
+   * Format token as you need. Such as:
+   *
+   * - rename token
+   * - merge token
+   * - delete token
+   *
+   * This should always be the same since it's one time process.
+   * It's ok to useMemo outside but this has better cache strategy.
+   */
+  formatToken?: (mergedToken: any) => DerivativeToken;
 }
 
 const tokenKeys = new Map<string, number>();
@@ -63,9 +74,9 @@ export default function useCacheToken<
 >(
   theme: Theme<any, any>,
   tokens: Partial<DesignToken>[],
-  option: Option = {},
+  option: Option<DerivativeToken> = {},
 ): [DerivativeToken & { _tokenKey: string }, string] {
-  const { salt = '', override = EMPTY_OVERRIDE } = option;
+  const { salt = '', override = EMPTY_OVERRIDE, formatToken } = option;
 
   // Basic - We do basic cache here
   const mergedToken = React.useMemo(
@@ -90,11 +101,16 @@ export default function useCacheToken<
       const derivativeToken = theme.getDerivativeToken(mergedToken);
 
       // Merge with override
-      const mergedDerivativeToken = {
+      let mergedDerivativeToken = {
         ...derivativeToken,
         ...override,
       };
 
+      // Format if needed
+      if (formatToken) {
+        mergedDerivativeToken = formatToken(mergedDerivativeToken);
+      }
+
       // Optimize for `useStyleRegister` performance
       const tokenKey = token2key(mergedDerivativeToken, salt);
       mergedDerivativeToken._tokenKey = tokenKey;
diff --git a/tests/index.spec.tsx b/tests/index.spec.tsx
index 689aa3f3..6e47e57b 100644
--- a/tests/index.spec.tsx
+++ b/tests/index.spec.tsx
@@ -212,11 +212,16 @@ describe('csssinjs', () => {
   });
 
   describe('override', () => {
-    const genStyle = (token: DerivativeToken): CSSInterpolation => ({
+    interface MyDerivativeToken extends DerivativeToken {
+      color: string;
+    }
+
+    const genStyle = (token: MyDerivativeToken): CSSInterpolation => ({
       '.box': {
         width: 93,
         lineHeight: 1,
         backgroundColor: token.primaryColor,
+        color: token.color,
       },
     });
 
@@ -226,8 +231,12 @@ describe('csssinjs', () => {
       propToken?: DesignToken;
       override: object;
     }) => {
-      const [token] = useCacheToken<DerivativeToken>(theme, [baseToken], {
+      const [token] = useCacheToken<MyDerivativeToken>(theme, [baseToken], {
         override,
+        formatToken: (origin: DerivativeToken) => ({
+          ...origin,
+          color: origin.primaryColor,
+        }),
       });
 
       useStyleRegister({ theme, token, path: ['.box'] }, () => [
@@ -255,6 +264,7 @@ describe('csssinjs', () => {
 
       const style = styles[0];
       expect(style.innerHTML).toContain('background-color:#010203;');
+      expect(style.innerHTML).toContain('color:#010203;');
 
       wrapper.unmount();
     });