feat: supprt complex key and initializer

baozouai · Apr 22, 2023 · be32932 · be32932
1 parent f2de710
commit be32932
Show file tree

Hide file tree

Showing 5 changed files with 202 additions and 89 deletions.
diff --git a/package.json b/package.json
@@ -38,6 +38,7 @@
     "prepare": "husky install"
   },
   "dependencies": {
+    "@babel/generator": "^7.21.4",
     "@babel/helper-plugin-utils": "^7.20.2",
     "@babel/plugin-syntax-typescript": "^7.21.4"
   },
@@ -49,6 +50,7 @@
     "@commitlint/cli": "^17.4.4",
     "@commitlint/config-conventional": "^17.4.4",
     "@types/babel__core": "^7.20.0",
+    "@types/babel__generator": "^7.6.4",
     "@types/babel__helper-plugin-utils": "^7.10.0",
     "@types/node": "^18.15.13",
     "@vitest/coverage-c8": "^0.30.1",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/babel-plugin-enum-to-object.ts b/src/babel-plugin-enum-to-object.ts
@@ -1,8 +1,9 @@
 import { declare } from '@babel/helper-plugin-utils'
 // @ts-ignore
 import syntaxTypeScript from '@babel/plugin-syntax-typescript'
-import type { BinaryExpression } from '@babel/types'
+import type { BinaryExpression, Identifier } from '@babel/types'
 import * as t from '@babel/types'
+import generater from '@babel/generator'
 interface BabelPluginEnumToObjectOptions {
   /**
    * need reflect ? default true
@@ -25,23 +26,25 @@ interface BabelPluginEnumToObjectOptions {
   reflect?: boolean
 }
 
-function getValueFromBinaryExpression(node: BinaryExpression, mapValue: Map<string | number, number | string>) {
+function getValueFromBinaryExpression(node: BinaryExpression, mapValue: Map<number | string | t.Identifier, number | string | t.Identifier>, unMapValue: (t.Identifier | t.MemberExpression)[] = []) {
   const { left, operator, right } = node
   let leftValue = 0
   let rightValue = 0
 
   if (t.isBinaryExpression(left))
-    leftValue = getValueFromBinaryExpression(left, mapValue)
+    leftValue = getValueFromBinaryExpression(left, mapValue, unMapValue)
   else if (t.isNumericLiteral(left))
     leftValue = left.value
   else if (t.isIdentifier(left))
     leftValue = +mapValue.get(left.name)!
+
   if (t.isBinaryExpression(right))
-    rightValue = getValueFromBinaryExpression(right, mapValue)
+    rightValue = getValueFromBinaryExpression(right, mapValue, unMapValue)
   else if (t.isNumericLiteral(right))
     rightValue = right.value
   else if (t.isIdentifier(right))
     rightValue = +mapValue.get(right.name)!
+
   if (operator === '+')
     return leftValue + rightValue
 
@@ -93,7 +96,7 @@ export default declare<BabelPluginEnumToObjectOptions>((api, options) => {
         const { node } = path
         const { id, members } = node
         let preNum = -1
-        const targetMap = new Map<string | number, number | string>()
+        const targetMap = new Map<number | string | t.Identifier, number | string | t.Identifier>()
 
         members.forEach((member) => {
           let { initializer, id: memberId } = member
@@ -110,22 +113,31 @@ export default declare<BabelPluginEnumToObjectOptions>((api, options) => {
           else
             key = memberId.value
 
-          let value: number | string = preNum
-          if (t.isStringLiteral(initializer))
+          let value: number | string | t.Identifier = preNum
+          if (t.isStringLiteral(initializer)) {
             value = initializer.value
-
-          else if (t.isBinaryExpression(initializer))
+          }
+          else if (t.isBinaryExpression(initializer)) {
             value = getValueFromBinaryExpression(initializer, targetMap)
-
+          }
+          else if (!t.isNumericLiteral(initializer)) {
+            const { code } = generater(initializer)
+            value = t.identifier(code)
+          }
           targetMap.set(key, value)
 
           if (reflect)
             targetMap.set(value, key)
         })
+
         const obj = t.variableDeclarator(
           id,
           t.objectExpression(
-            [...targetMap.entries()].map(([key, value]) => t.objectProperty(t.identifier(String(key)), typeof value === 'string' ? t.stringLiteral(value) : t.numericLiteral(value))),
+            [...targetMap.entries()].map(([key, value]) => {
+              const objectKey = t.isIdentifier(key as Identifier) ? t.identifier(`[${(key as Identifier).name}]`) : t.stringLiteral(String(key))
+              const objectValue = t.isIdentifier(value as Identifier) ? (value as Identifier) : typeof value === 'string' ? t.stringLiteral(value) : t.numericLiteral(value as number)
+              return t.objectProperty(objectKey, objectValue)
+            }),
           ),
         )
         const constObjVariable = t.variableDeclaration('const', [obj])

diff --git a/test/__snapshots__/index.test.ts.snap b/test/__snapshots__/index.test.ts.snap
@@ -2,118 +2,171 @@
 
 exports[`Transforms NumericLiteral enum  with reflect false 1`] = `
 "export const Direction = {
-  Up: 0,
-  Down: 1,
-  Left: 2,
-  Right: 3
+  'Up': 0,
+  'Down': 1,
+  'Left': 2,
+  'Right': 3
 };"
 `;
 
 exports[`Transforms NumericLiteral enum 1`] = `
 "export const Direction = {
-  Up: 0,
-  0: 'Up',
-  Down: 1,
-  1: 'Down',
-  Left: 2,
-  2: 'Left',
-  Right: 3,
-  3: 'Right'
+  'Up': 0,
+  '0': 'Up',
+  'Down': 1,
+  '1': 'Down',
+  'Left': 2,
+  '2': 'Left',
+  'Right': 3,
+  '3': 'Right'
 };"
 `;
 
 exports[`Transforms calc enum  with reflect false 1`] = `
 "export const MyEnum = {
-  A: 0,
-  B: 1,
-  C: 20,
-  D: 21,
-  E: 4200,
-  F: 4199,
-  G: 4202,
-  H: 2101,
-  I: 5,
-  J: 74088000000,
-  K: 16800,
-  L: 1050,
-  M: 2100,
-  N: 1,
-  O: 4203,
-  P: 4202
+  'A': 0,
+  'B': 1,
+  'C': 20,
+  'D': 21,
+  'E': 4200,
+  'F': 4199,
+  'G': 4202,
+  'H': 2101,
+  'I': 5,
+  'J': 74088000000,
+  'K': 16800,
+  'L': 1050,
+  'M': 2100,
+  'N': 1,
+  'O': 4203,
+  'P': 4202
 };"
 `;
 
 exports[`Transforms calc enum 1`] = `
 "export const MyEnum = {
-  A: 0,
-  0: 'A',
-  B: 1,
-  1: 'N',
-  C: 20,
-  20: 'C',
-  D: 21,
-  21: 'D',
-  E: 4200,
-  4200: 'E',
-  F: 4199,
-  4199: 'F',
-  G: 4202,
-  4202: 'P',
-  H: 2101,
-  2101: 'H',
-  I: 5,
-  5: 'I',
-  J: 74088000000,
-  74088000000: 'J',
-  K: 16800,
-  16800: 'K',
-  L: 1050,
-  1050: 'L',
-  M: 2100,
-  2100: 'M',
-  N: 1,
-  O: 4203,
-  4203: 'O',
-  P: 4202
+  'A': 0,
+  '0': 'A',
+  'B': 1,
+  '1': 'N',
+  'C': 20,
+  '20': 'C',
+  'D': 21,
+  '21': 'D',
+  'E': 4200,
+  '4200': 'E',
+  'F': 4199,
+  '4199': 'F',
+  'G': 4202,
+  '4202': 'P',
+  'H': 2101,
+  '2101': 'H',
+  'I': 5,
+  '5': 'I',
+  'J': 74088000000,
+  '74088000000': 'J',
+  'K': 16800,
+  '16800': 'K',
+  'L': 1050,
+  '1050': 'L',
+  'M': 2100,
+  '2100': 'M',
+  'N': 1,
+  'O': 4203,
+  '4203': 'O',
+  'P': 4202
+};"
+`;
+
+exports[`Transforms complex enum  with reflect false 1`] = `
+"export const Other = {
+  'A': 4,
+  'B': 5,
+  'C': 6
+};
+export const MyEnum = {
+  'GP-2120TU': Other.A,
+  'GP-M32B': Other.B,
+  'GP-R222C': Other.C
+};"
+`;
+
+exports[`Transforms complex enum 1`] = `
+"export const Other = {
+  'A': 4,
+  '4': 'A',
+  'B': 5,
+  '5': 'B',
+  'C': 6,
+  '6': 'C'
+};
+export const MyEnum = {
+  'GP-2120TU': Other.A,
+  [Other.A]: 'GP-2120TU',
+  'GP-M32B': Other.B,
+  [Other.B]: 'GP-M32B',
+  'GP-R222C': Other.C,
+  [Other.C]: 'GP-R222C'
+};"
+`;
+
+exports[`Transforms initByOtherIdentifiter enum 1`] = `
+"const WW = 2;
+export const MyEnum = {
+  'GP-2120TU': WW,
+  [WW]: 'GP-2120TU',
+  'GP-M32B': 3,
+  '3': 'GP-M32B',
+  'GP-R222C': 4,
+  '4': 'GP-R222C'
+};"
+`;
+
+exports[`Transforms initByOtherIdentifiter enum false 1`] = `
+"const WW = 2;
+export const MyEnum = {
+  'GP-2120TU': WW,
+  'GP-M32B': 3,
+  'GP-R222C': 4
 };"
 `;
 
 exports[`Transforms string enum  with reflect false 1`] = `
 "export const Direction = {
-  Up: 'Up',
-  Down: 'Down',
-  Left: 'Left',
-  Right: 'Right'
+  'Up': 'Up',
+  'Down': 'Down',
+  'Left': 'Left',
+  'Right': 'Right'
 };"
 `;
 
 exports[`Transforms string enum 1`] = `
 "export const Direction = {
-  Up: 'Up',
-  Down: 'Down',
-  Left: 'Left',
-  Right: 'Right'
+  'Up': 'Up',
+  'Down': 'Down',
+  'Left': 'Left',
+  'Right': 'Right'
 };"
 `;
 
 exports[`Transforms string number mix  with reflect false 1`] = `
 "export const Status = {
-  PAID: 0,
-  UNPAID: '21',
-  PART: 44,
-  HALF: 45
+  'PAID': 0,
+  'UNPAID': '21',
+  'PART': 44,
+  'HALF': 45
 };"
 `;
 
 exports[`Transforms string number mix 1`] = `
 "export const Status = {
-  PAID: 0,
-  0: 'PAID',
-  UNPAID: '21',
-  21: 'UNPAID',
-  PART: 44,
-  44: 'PART',
-  HALF: 45,
-  45: 'HALF'
+  'PAID': 0,
+  '0': 'PAID',
+  'UNPAID': '21',
+  '21': 'UNPAID',
+  'PART': 44,
+  '44': 'PART',
+  'HALF': 45,
+  '45': 'HALF'
 };"
 `;