From 304530d2106ac0575fe06cc8a532be48e441d924 Mon Sep 17 00:00:00 2001
From: allohamora <herman30698@gmail.com>
Date: Thu, 27 Jul 2023 20:43:16 +0300
Subject: [PATCH] feat: improve defaultFor with rest return value

---
 __tests__/unit/env.picker.spec.ts | 5 ++++-
 src/env/env.picker.ts             | 8 +++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/__tests__/unit/env.picker.spec.ts b/__tests__/unit/env.picker.spec.ts
index 09dbd2d..1f462d5 100644
--- a/__tests__/unit/env.picker.spec.ts
+++ b/__tests__/unit/env.picker.spec.ts
@@ -40,7 +40,7 @@ describe('EnvPicker', () => {
       ).toBe(undefined);
 
       expect(new EnvPicker(false as boolean | undefined, 'test').defaultFor({ test: true }).value()).toBe(false);
-      expect(new EnvPicker('', 'test').defaultFor({ test: '123' }).value()).toBe('');
+      expect(new EnvPicker('' as string | undefined, 'test').defaultFor({ test: '123' }).value()).toBe('');
     });
   });
 
@@ -48,6 +48,9 @@ describe('EnvPicker', () => {
     it('maps state', () => {
       expect(new EnvPicker('123', 'test').map(Number).value()).toBe(123);
       expect(new EnvPicker(undefined as string | undefined, 'test').default('123').map(Number).value()).toBe(123);
+      expect(
+        new EnvPicker(undefined as string | undefined, 'test').defaultFor({ rest: '123' }).map(Number).value(),
+      ).toBe(123);
       expect(new EnvPicker('false', 'test').default('123').map(JSON.parse).value()).toBe(false);
       expect(new EnvPicker(undefined as string | undefined, 'test').map((state) => !!state).value()).toBe(false);
     });
diff --git a/src/env/env.picker.ts b/src/env/env.picker.ts
index d6abcfe..4ece5da 100644
--- a/src/env/env.picker.ts
+++ b/src/env/env.picker.ts
@@ -31,10 +31,12 @@ export class EnvPicker<NodeEnv extends string, State> {
     return this as unknown as EnvPicker<NodeEnv, NewState | NonNullable<State>>;
   }
 
-  public defaultFor(envRecord: EnvRecord<NodeEnv, State>) {
-    this.state ??= envRecord[this.nodeEnv] ?? (envRecord.rest as State);
+  public defaultFor<NewState extends State, EnvOptions extends EnvRecord<NodeEnv, NewState>>(envRecord: EnvOptions) {
+    this.state ??= (envRecord[this.nodeEnv] ?? envRecord.rest) as NewState;
 
-    return this;
+    return this as unknown as EnvOptions extends { rest: NewState }
+      ? EnvPicker<NodeEnv, NonNullable<NewState>>
+      : EnvPicker<NodeEnv, NewState>;
   }
 
   public map<Result>(mapper: (state: State) => Result) {