From 937a3f163ff46538e01ad0a24a9e5af3798844bf Mon Sep 17 00:00:00 2001
From: pandaniell <daniell.wdb@icloud.com>
Date: Fri, 5 Mar 2021 20:16:21 +0100
Subject: [PATCH] feat: allow union type generation

---
 README.md        | 21 ++++++++++++++++++---
 gen-env-types.js | 12 +++++++++++-
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 0170799..cfaf446 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@ STRIPE_ACCESS_TOKEN=qoi120wqe
 
 And generates a `.d.ts` file
 
-```
+```typescript
 declare namespace NodeJS {
   export interface ProcessEnv {
     SESSION_SECRET: string;
@@ -20,13 +20,28 @@ declare namespace NodeJS {
 
 Now `process.env.SESSION_SECRET` will autocomplete and be type-safe.
 
-## Usage
+If you want to generate a union instead of string type, add an inline comment to your `.env` file:
 
 ```
+NODE_ENV = "production" # production | development
+```
+
+```typescript
+declare namespace NodeJS {
+  export interface ProcessEnv {
+    NODE_ENV: "production" | "development";
+  }
+}
+```
+
+## Usage
+
+```bash
 npx gen-env-types path/to/.env
 ```
 
 ## Options
+
 ```
   -V, --version               Show version number
   -h, --help                  Show usage information
@@ -36,6 +51,6 @@ npx gen-env-types path/to/.env
 
 ## Example with options
 
-```
+```bash
 npx gen-env-types .env -o src/types/env.d.ts -e .
 ```
diff --git a/gen-env-types.js b/gen-env-types.js
index cc990f5..2c456f1 100644
--- a/gen-env-types.js
+++ b/gen-env-types.js
@@ -115,7 +115,17 @@ function writeEnvTypes(envString, path) {
     ${envString
       .split("\n")
       .filter((line) => line.trim() && line.trim().indexOf("#") !== 0)
-      .map((x, i) => `${i ? "    " : ""}${x.trim().split("=")[0]}: string;`)
+      .map((x, i) => {
+        const union = x.split('"').pop().trim()
+
+        if (union) {
+          const stringLiterals = union.replace("#", "").split(" | ")
+
+          return `${i ? "    " : ""}${x.split("=")[0]}: ${stringLiterals.map((text) => `"${text.trim()}"`).join(" | ")};`
+        }
+
+        return `${i ? "    " : ""}${x.trim().split("=")[0]}: string;`
+      })
       .join("\n")}
   }
 }