From 80c7f7c1057c473357f801daec9489dfeee50729 Mon Sep 17 00:00:00 2001
From: Victor Bo <10667379@qq.com>
Date: Mon, 29 Apr 2024 14:04:21 +0800
Subject: [PATCH] feat: support custom types

---
 playground/types/{auto-env.dts => auto-env.d.ts} |  2 +-
 playground/vite.config.ts                        |  6 +++++-
 src/core/context.ts                              |  4 ++--
 src/core/env.ts                                  | 12 +++++++++---
 src/core/options.ts                              |  1 +
 src/types.ts                                     |  6 ++++++
 6 files changed, 24 insertions(+), 7 deletions(-)
 rename playground/types/{auto-env.dts => auto-env.d.ts} (93%)

diff --git a/playground/types/auto-env.dts b/playground/types/auto-env.d.ts
similarity index 93%
rename from playground/types/auto-env.dts
rename to playground/types/auto-env.d.ts
index d3823ad..694a73f 100644
--- a/playground/types/auto-env.dts
+++ b/playground/types/auto-env.d.ts
@@ -16,7 +16,7 @@ interface ImportMetaEnv {
    */
   readonly VITE_GE: string
   readonly VITE_OBJECT: string
-  readonly VITE_STRING: string
+  readonly VITE_STRING: 'abc' | 'bc'
   readonly VITE_ABC: number
 }
 
diff --git a/playground/vite.config.ts b/playground/vite.config.ts
index 0867173..f462899 100644
--- a/playground/vite.config.ts
+++ b/playground/vite.config.ts
@@ -6,7 +6,11 @@ export default defineConfig({
   plugins: [
     Inspect(),
     VitePlugin({
-      dts: 'types/auto-env.dts',
+      dts: 'types/auto-env.d.ts',
+      custom: {
+        VITE_GE: 'string',
+        VITE_STRING: `'abc' | 'bc'`,
+      },
     }),
   ],
 })
diff --git a/src/core/context.ts b/src/core/context.ts
index 2dbd639..84e4a81 100644
--- a/src/core/context.ts
+++ b/src/core/context.ts
@@ -26,7 +26,7 @@ export class Context {
   }
 
   scanEnv() {
-    const { dts, includes } = this.options
+    const { dts, includes, custom } = this.options
     if (dts !== false) {
       const envFiles = fg.sync(
         includes,
@@ -41,7 +41,7 @@ export class Context {
 
       envFiles.forEach((path) => {
         const content = readFileSync(path, 'utf-8')
-        const envMap = parseMetaEnv(content)
+        const envMap = parseMetaEnv(content, custom)
         envMap.forEach((env) => {
           this._env.set(env.label, { ...env })
         })
diff --git a/src/core/env.ts b/src/core/env.ts
index 9441b18..bbd4679 100644
--- a/src/core/env.ts
+++ b/src/core/env.ts
@@ -1,6 +1,7 @@
 import { existsSync, mkdirSync, writeFileSync } from 'node:fs'
 import { dirname, resolve } from 'node:path'
 import createDebugger from 'debug'
+import { isKeyOfObj } from '@vtrbo/utils'
 import type { ResolvedOptions } from '../types'
 import { getLikeType } from './utils'
 import { VITE_PLUGIN_NAME } from './constant'
@@ -11,11 +12,11 @@ export interface Env {
   remark: string
   label: string
   value: string
-  likely: 'string' | 'boolean' | 'number'
+  likely: 'string' | 'boolean' | 'number' | string
   required: boolean
 }
 
-export function parseMetaEnv(data: string): Env[] {
+export function parseMetaEnv(data: string, custom: ResolvedOptions['custom']): Env[] {
   data = data.replace(/\r\n?/gm, '\n')
   const regexp = /(?:^|^)\s*((?:\s*#.+\n)*)?(?:export\s+)?([\w.-]+)(?:\s*=\s*?|:\s+?)(\s*'(?:\\'|[^'])*'|\s*"(?:\\"|[^"])*"|\s*`(?:\\`|[^`])*`|[^#\r\n]+)?[^#\r\n]*(#.*)?(?:$|$)/gm
 
@@ -32,13 +33,18 @@ export function parseMetaEnv(data: string): Env[] {
     if (isString)
       value = value.replace(quoteRegExp, '$2').replace(/\\n/g, '\n').replace(/\\r/g, '\r')
 
+    let likely = isString ? 'string' : getLikeType(value)
+
+    if (isKeyOfObj(custom, label))
+      likely = custom[label]
+
     const required = !['', null, undefined].includes(value)
 
     meteEnv.push({
       remark,
       label,
       value,
-      likely: isString ? 'string' : getLikeType(value),
+      likely,
       required,
     })
   }
diff --git a/src/core/options.ts b/src/core/options.ts
index f878a87..3ad1e7e 100644
--- a/src/core/options.ts
+++ b/src/core/options.ts
@@ -7,5 +7,6 @@ export function resolveOptions(options?: Options): ResolvedOptions {
     includes: options?.includes ? toArray(options.includes) : ENV_INCLUDES,
     prefix: options?.prefix ? toArray(options.prefix) : ENV_PREFIX,
     dts: options?.dts ?? ENV_DTS,
+    custom: options?.custom ? options.custom : {},
   }
 }
diff --git a/src/types.ts b/src/types.ts
index 36be8b8..6edb238 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -17,10 +17,16 @@ export interface Options {
    * dts file path
    */
   dts?: string | false
+
+  /**
+   * 自定义类型
+   */
+  custom?: Record<string, string>
 }
 
 export interface ResolvedOptions {
   includes: string[]
   prefix: string[]
   dts: string | false
+  custom: Record<string, string>
 }