feat: support custom getComputedToken (#135)

* feat: support custom getComputedToken * test: add test * chore: add comment * chore: more comment * chore: update test case
ant-design · Jul 26, 2023 · c0f9c20 · c0f9c20
1 parent f784089
commit c0f9c20
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 5 deletions.
diff --git a/src/hooks/useCacheToken.tsx b/src/hooks/useCacheToken.tsx
@@ -15,7 +15,7 @@ const hashPrefix =
     ? 'css-dev-only-do-not-override'
     : 'css';
 
-export interface Option<DerivativeToken> {
+export interface Option<DerivativeToken, DesignToken> {
   /**
    * Generate token with salt.
    * This is used to generate different hashId even same derivative token for different version.
@@ -33,6 +33,14 @@ export interface Option<DerivativeToken> {
    * It's ok to useMemo outside but this has better cache strategy.
    */
   formatToken?: (mergedToken: any) => DerivativeToken;
+  /**
+   * Get final token with origin token, override token and theme.
+   * The parameters do not contain formatToken since it's passed by user.
+   * @param origin The original token.
+   * @param override Extra tokens to override.
+   * @param theme Theme instance. Could get derivative token by `theme.getDerivativeToken`
+   */
+  getComputedToken?: (origin: DesignToken, override: object, theme: Theme<any, any>) => DerivativeToken;
 }
 
 const tokenKeys = new Map<string, number>();
@@ -112,12 +120,17 @@ export default function useCacheToken<
 >(
   theme: Theme<any, any>,
   tokens: Partial<DesignToken>[],
-  option: Option<DerivativeToken> = {},
+  option: Option<DerivativeToken, DesignToken> = {},
 ): [DerivativeToken & { _tokenKey: string }, string] {
   const {
     cache: { instanceId },
   } = useContext(StyleContext);
-  const { salt = '', override = EMPTY_OVERRIDE, formatToken } = option;
+  const {
+    salt = '',
+    override = EMPTY_OVERRIDE,
+    formatToken,
+    getComputedToken: compute
+  } = option;
 
   // Basic - We do basic cache here
   const mergedToken = React.useMemo(
@@ -139,7 +152,7 @@ export default function useCacheToken<
     'token',
     [salt, theme.id, tokenStr, overrideTokenStr],
     () => {
-      const mergedDerivativeToken = getComputedToken(
+      const mergedDerivativeToken = compute ? compute(mergedToken, override, theme) : getComputedToken(
         mergedToken,
         override,
         theme,

diff --git a/src/util.ts b/src/util.ts
@@ -1,13 +1,16 @@
 import hash from '@emotion/hash';
 import canUseDom from 'rc-util/lib/Dom/canUseDom';
 import { removeCSS, updateCSS } from 'rc-util/lib/Dom/dynamicCSS';
+import { Theme } from './theme';
 
 export function flattenToken(token: any) {
   let str = '';
   Object.keys(token).forEach((key) => {
     const value = token[key];
     str += key;
-    if (value && typeof value === 'object') {
+    if (value instanceof Theme) {
+      str += value.id;
+    } else if (value && typeof value === 'object') {
       str += flattenToken(value);
     } else {
       str += value;

diff --git a/tests/index.spec.tsx b/tests/index.spec.tsx
@@ -8,8 +8,10 @@ import {
   Theme,
   useCacheToken,
   useStyleRegister,
+  createTheme
 } from '../src';
 import { ATTR_MARK, ATTR_TOKEN, CSS_IN_JS_INSTANCE } from '../src/StyleContext';
+import type { DerivativeFunc } from '../src';
 
 interface DesignToken {
   primaryColor: string;
@@ -551,4 +553,59 @@ describe('csssinjs', () => {
 
     render(<Demo />);
   });
+
+  it('should support custom getComputedToken', () => {
+    const genDemoStyle = (token: any): CSSInterpolation => ({
+      div: {
+        color: token.myToken,
+        background: token.primaryColor,
+      },
+    });
+
+    const Demo = ({myToken, theme: customTheme}: { myToken?: string, theme?: DerivativeFunc<any, any> }) => {
+      const [token, hashId] = useCacheToken<DerivativeToken>(theme, [{primaryColor: 'blue'}], {
+        salt: 'test',
+        override: {
+          myToken,
+          theme: customTheme && createTheme(customTheme)
+        },
+        getComputedToken: (origin, override: any, myTheme) => {
+          const mergedToken = myTheme.getDerivativeToken(origin);
+          return {
+            ...mergedToken,
+            myToken: override.myToken,
+            ...(override.theme?.getDerivativeToken(mergedToken) ?? {}),
+          }
+        }
+      });
+
+      useStyleRegister(
+        { theme, token, hashId, path: ['cssinjs-getComputedToken'] },
+        () => [genDemoStyle(token)],
+      );
+
+      return <div className={classNames('box', hashId)} />;
+    };
+
+    const { rerender } =render(<Demo myToken="test" />);
+
+    const styles = Array.from(document.head.querySelectorAll('style'));
+    expect(styles).toHaveLength(1);
+    expect(styles[0].innerHTML).toContain('color:test');
+    expect(styles[0].innerHTML).toContain('background:blue');
+
+    rerender(<Demo myToken="apple" />);
+
+    const styles2 = Array.from(document.head.querySelectorAll('style'));
+    expect(styles2).toHaveLength(1);
+    expect(styles2[0].innerHTML).toContain('color:apple');
+    expect(styles2[0].innerHTML).toContain('background:blue');
+
+    rerender(<Demo myToken="banana" theme={(origin) => ({...origin, primaryColor: 'green'})} />);
+
+    const styles3 = Array.from(document.head.querySelectorAll('style'));
+    expect(styles3).toHaveLength(1);
+    expect(styles3[0].innerHTML).toContain('color:banana');
+    expect(styles3[0].innerHTML).toContain('background:green');
+  })
 });