Merge pull request #96 from Nikaple/feat/custom-key-serializer

Nikaple · Mar 22, 2024 · 2bd805b · 2bd805b
2 parents f2104c9 + 9ab90f5
commit 2bd805b
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 19 deletions.
diff --git a/src/flatten.test.ts b/src/flatten.test.ts
@@ -203,6 +203,34 @@ describe('flattenObject', () => {
     testFlatten({ cat }, { cat });
   });
 
+  it('should be able to custom key serializer', () => {
+    expect(
+      flatten(
+        { a: [0, { b: [1], c: { d: [2], 5: [6] } }] },
+        {
+          serializeFlattenKey(key, prefix, meta) {
+            if (meta.isArrayIndex) {
+              return `${prefix}.[${key}]`;
+            }
+            if (
+              meta.hasSpecialCharacters ||
+              meta.isEmpty ||
+              /^\d+$/.test(key)
+            ) {
+              return `${prefix}[${JSON.stringify(key)}]`;
+            }
+            return prefix ? `${prefix}.${key}` : key;
+          },
+        },
+      ),
+    ).toEqual({
+      'a.[0]': 0,
+      'a.[1].b.[0]': 1,
+      'a.[1].c.d.[0]': 2,
+      'a.[1].c["5"].[0]': 6,
+    });
+  });
+
   it('should not throw on illegal input', () => {
     expect(flatten('' as any)).toEqual({});
   });

diff --git a/src/flatten.ts b/src/flatten.ts
@@ -1,5 +1,6 @@
 import { deepSet } from './deep-set';
 import {
+  config,
   extractCircularKey,
   formatCircularKey,
   isObject,
@@ -23,22 +24,7 @@ export const flatten = <T>(
   options?: UniFlattenOptions,
 ): Record<string, T> => {
   const seen = new Map();
-  const getKey = (key: string, prefix: string, isNumber: boolean) => {
-    let k;
-    if (
-      /[.'"\s\\\b\f\n\r\t\v{}()[\];,<>=!+\-*%&|^~?:]|^\d+\D/.test(key) ||
-      key === ''
-    ) {
-      // use brackets if key contains special characters
-      k = `[${JSON.stringify(key)}]`;
-    } else if (/^\d+$/.test(key)) {
-      // use [0] for arrays, and ["0"] for numeric keys
-      k = isNumber ? `[${key}]` : `["${key}"]`;
-    } else {
-      k = key;
-    }
-    return prefix ? `${prefix}${/^\[/.test(k) ? '' : '.'}${k}` : k;
-  };
+  const serializer = options?.serializeFlattenKey || config.serializeFlattenKey;
   const helper = (obj: any, prefix: string, result: any = {}) => {
     if (!isObject(obj)) {
       return result;
@@ -53,7 +39,11 @@ export const flatten = <T>(
     // recursively handle arrays
     if (Array.isArray(obj)) {
       obj.forEach((item, i) => {
-        const key = getKey(String(i), prefix, true);
+        const key = serializer(String(i), prefix, {
+          isArrayIndex: true,
+          isEmpty: false,
+          hasSpecialCharacters: false,
+        });
         if (isObject(item)) {
           const res = helper(item, key, result);
           return res;
@@ -66,7 +56,12 @@ export const flatten = <T>(
     // recursively handle plain objects
     if (isPlainObject(obj)) {
       Object.entries(obj).forEach(([k, item]) => {
-        const key = getKey(k, prefix, false);
+        const key = serializer(k, prefix, {
+          isEmpty: k === '',
+          isArrayIndex: false,
+          hasSpecialCharacters:
+            /[.'"\s\\\b\f\n\r\t\v{}()[\];,<>=!+\-*%&|^~?:]|^\d+\D/.test(k),
+        });
         if (isObject(item)) {
           const res = helper(item, key, result);
           return res;

diff --git a/src/internal.ts b/src/internal.ts
@@ -7,7 +7,24 @@ export const SPECIAL_CHARACTER_REGEX =
 export const config = {
   strict: false,
   circularReference: 'string' as const,
-};
+  serializeFlattenKey: (
+    key: string,
+    prefix: string,
+    meta: {
+      isArrayIndex: boolean;
+      isEmpty: boolean;
+      hasSpecialCharacters: boolean;
+    },
+  ) => {
+    if (meta.isArrayIndex) {
+      return `${prefix}[${key}]`;
+    }
+    if (meta.hasSpecialCharacters || meta.isEmpty || /^\d+$/.test(key)) {
+      return `${prefix}[${JSON.stringify(key)}]`;
+    }
+    return prefix ? `${prefix}.${key}` : key;
+  },
+} satisfies UniFlattenOptions;
 
 export const mergeConfig = (options?: UniFlattenOptions) => ({
   ...config,

diff --git a/src/type.ts b/src/type.ts
@@ -1,4 +1,14 @@
 export interface UniFlattenOptions {
   circularReference?: 'string' | 'symbol' | 'null';
   strict?: boolean;
+  serializeFlattenKey?: (
+    /** current key to serialize */
+    key: string,
+    prefix: string,
+    meta: {
+      isArrayIndex: boolean;
+      isEmpty: boolean;
+      hasSpecialCharacters: boolean;
+    },
+  ) => string;
 }