fix: lens extensions (#49)

Author: aspirisen

.github/workflows/ci.yml

@@ -23,6 +23,19 @@ jobs:
       - name: Prettier
         run: bun prettier
 
+  typecheck:
+    name: Typecheck
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Setup environment
+        uses: ./.github/actions/setup
+
+      - name: Run typecheck
+        run: bun typecheck
+
   test-unit:
     name: Test Unit
     runs-on: ubuntu-latest

package.json

@@ -54,7 +54,8 @@
     "prettier:fix": "prettier --write '**/*.md'",
     "storybook": "storybook dev -p 6006",
     "typecheck": "tsc --noEmit",
-    "test": "vitest watch",
+    "test": "vitest",
+    "test:watch": "vitest watch",
     "test:e2e": "vitest run --project=e2e --coverage",
     "test:unit": "vitest run --project=unit --coverage"
   },

src/LensCore.ts

@@ -33,7 +33,7 @@ export class LensCore<T extends FieldValues> {
     control: Control<TFieldValues>,
     cache?: LensesStorage<TFieldValues>,
   ): Lens<TFieldValues> {
-    return new LensCore(control, '', cache) as unknown as Lens<TFieldValues>;
+    return new this(control, '', cache) as unknown as Lens<TFieldValues>;
   }
 
   public focus(prop: string | number): LensCore<T> {
@@ -57,7 +57,7 @@ export class LensCore<T extends FieldValues> {
 
     if (Array.isArray(this.override)) {
       const [template] = this.override;
-      const result = new LensCore(this.control, nestedPath, this.cache);
+      const result = new (this.constructor as typeof LensCore)(this.control, nestedPath, this.cache);
       result.isArrayItemReflection = true;
       result.override = template;
 
@@ -68,14 +68,14 @@ export class LensCore<T extends FieldValues> {
       const overriddenLens: LensCore<T> | undefined = get(this.override, propString);
 
       if (!overriddenLens) {
-        const result = new LensCore(this.control, nestedPath, this.cache);
+        const result = new (this.constructor as typeof LensCore)(this.control, nestedPath, this.cache);
         this.cache?.set(result, nestedPath);
         return result;
       }
 
       if (this.isArrayItemReflection) {
         const arrayItemNestedPath = `${this.path}.${overriddenLens.path}`;
-        const result = new LensCore(this.control, arrayItemNestedPath, this.cache);
+        const result = new (this.constructor as typeof LensCore)(this.control, arrayItemNestedPath, this.cache);
         this.cache?.set(result, arrayItemNestedPath);
         return result;
       } else {
@@ -84,7 +84,7 @@ export class LensCore<T extends FieldValues> {
       }
     }
 
-    const result = new LensCore(this.control, nestedPath, this.cache);
+    const result = new (this.constructor as typeof LensCore)(this.control, nestedPath, this.cache);
     this.cache?.set(result, nestedPath);
     return result;
   }
@@ -102,7 +102,7 @@ export class LensCore<T extends FieldValues> {
     }
 
     const nestedCache = new LensesStorage(this.control);
-    const template = new LensCore(this.control, this.path, nestedCache);
+    const template = new (this.constructor as typeof LensCore)(this.control, this.path, nestedCache);
 
     const dictionary = new Proxy(
       {},
@@ -120,7 +120,7 @@ export class LensCore<T extends FieldValues> {
     const override = getter(dictionary, template);
 
     if (Array.isArray(override)) {
-      const result = new LensCore(this.control, this.path, this.cache);
+      const result = new (this.constructor as typeof LensCore)(this.control, this.path, this.cache);
       template.path = '';
       result.override = getter(dictionary, template);
       result.reflectedKey = getter;

tests/lens-extension.test.ts

@@ -0,0 +1,125 @@
+import type { FieldValues } from 'react-hook-form';
+import { useForm } from 'react-hook-form';
+import { LensCore, LensesStorage } from '@hookform/lenses';
+import { renderHook } from '@testing-library/react';
+
+// Create a custom LensCore class that extends the base class
+class CustomLensCore<T extends FieldValues> extends LensCore<T> {
+  // Add a custom method
+  public getValue(): string {
+    return `Custom value at path: ${this.path}`;
+  }
+
+  // Add another custom method
+  public getCustomInfo(): string {
+    return `CustomLensCore instance at ${this.path}`;
+  }
+}
+
+test('LensCore extension should preserve custom methods when using focus()', () => {
+  const { result } = renderHook(() =>
+    useForm<{
+      user: {
+        name: string;
+        email: string;
+      };
+    }>({
+      defaultValues: {
+        user: {
+          name: 'John',
+          email: '[email protected]',
+        },
+      },
+    }),
+  );
+
+  const control = result.current.control;
+  const cache = new LensesStorage(control);
+  const rootLens = new CustomLensCore(control, '', cache);
+
+  // Focus on nested paths
+  const userLens = rootLens.focus('user');
+  const nameLens = userLens.focus('name');
+
+  // Verify that the focused lenses are instances of CustomLensCore
+  expect(userLens).toBeInstanceOf(CustomLensCore);
+  expect(nameLens).toBeInstanceOf(CustomLensCore);
+
+  // Verify that custom methods are available
+  expect((userLens as CustomLensCore<any>).getValue()).toBe('Custom value at path: user');
+  expect((userLens as CustomLensCore<any>).getCustomInfo()).toBe('CustomLensCore instance at user');
+
+  expect((nameLens as CustomLensCore<any>).getValue()).toBe('Custom value at path: user.name');
+  expect((nameLens as CustomLensCore<any>).getCustomInfo()).toBe('CustomLensCore instance at user.name');
+});
+
+test('LensCore extension should preserve custom methods when using reflect()', () => {
+  const { result } = renderHook(() =>
+    useForm<{
+      profile: {
+        contact: {
+          firstName: string;
+          phone: string;
+        };
+      };
+    }>({
+      defaultValues: {
+        profile: {
+          contact: {
+            firstName: 'Jane',
+            phone: '123-456-7890',
+          },
+        },
+      },
+    }),
+  );
+
+  const control = result.current.control;
+  const cache = new LensesStorage(control);
+  const rootLens = new CustomLensCore(control, '', cache);
+
+  // Use reflect to restructure
+  const contactLens = rootLens.reflect((_dictionary, l) => ({
+    name: l.focus('profile.contact.firstName'),
+    phoneNumber: l.focus('profile.contact.phone'),
+  }));
+
+  // Verify that the reflected lens is an instance of CustomLensCore
+  expect(contactLens).toBeInstanceOf(CustomLensCore);
+  expect((contactLens as CustomLensCore<any>).getValue()).toBe('Custom value at path: ');
+  expect((contactLens as CustomLensCore<any>).getCustomInfo()).toBe('CustomLensCore instance at ');
+
+  // Focus on the reflected structure
+  const nameLens = contactLens.focus('name');
+  expect(nameLens).toBeInstanceOf(CustomLensCore);
+});
+
+test('LensCore extension should preserve custom methods when using static create()', () => {
+  const { result } = renderHook(() =>
+    useForm<{
+      firstName: string;
+      lastName: string;
+    }>({
+      defaultValues: {
+        firstName: 'Test',
+        lastName: 'User',
+      },
+    }),
+  );
+
+  const control = result.current.control;
+  const cache = new LensesStorage(control);
+
+  // Use the static create method on the custom class
+  const lens = CustomLensCore.create(control, cache);
+
+  // Verify that it creates an instance of CustomLensCore
+  expect(lens).toBeInstanceOf(CustomLensCore);
+  expect((lens as unknown as CustomLensCore<any>).getValue()).toBe('Custom value at path: ');
+  expect((lens as unknown as CustomLensCore<any>).getCustomInfo()).toBe('CustomLensCore instance at ');
+
+  // Verify that focused lenses also preserve the custom class
+  const firstNameLens = lens.focus('firstName');
+  expect(firstNameLens).toBeInstanceOf(CustomLensCore);
+  expect((firstNameLens as unknown as CustomLensCore<any>).getValue()).toBe('Custom value at path: firstName');
+});

tests/types.test-d.ts

@@ -5,7 +5,7 @@ import { zodResolver } from '@hookform/resolvers/zod';
 import { expectTypeOf } from 'vitest';
 import { z } from 'zod';
 
-test('lens should use output type of the Controller', () => {
+test('lens should use input type of the Controller', () => {
   const schema = z.object({
     id: z.number().transform((val) => val.toString()),
   });
@@ -21,7 +21,7 @@ test('lens should use output type of the Controller', () => {
   });
 
   const lens = useLens({ control: form.control });
-  expectTypeOf<typeof lens>().toEqualTypeOf<Lens<Output>>();
+  expectTypeOf<typeof lens>().toEqualTypeOf<Lens<Input>>();
 });
 
 test('lens unwraps to the original type', () => {
@@ -214,28 +214,3 @@ test('generic assert for primitive union', () => {
   maybe.assert<string>();
   needsString(maybe); // should compile when assert works
 });
-
-test('should be able to narrow branded types', () => {
-  const EmailAddress = z.email().brand('EmailAddress');
-  type EmailAddress = z.infer<typeof EmailAddress>;
-
-  const FormSchema = z.object({ email: EmailAddress });
-  type FormValue = z.infer<typeof FormSchema>;
-  const { control } = useForm({ resolver: zodResolver(FormSchema) });
-  const lens = useLens({ control });
-
-  function check({ lens }: { lens: Lens<FormValue> }) {
-    return lens;
-  }
-
-  check({ lens });
-
-  const email = lens.focus('email');
-  expectTypeOf<typeof email>().toEqualTypeOf<Lens<string & z.core.$brand<'EmailAddress'>>>();
-
-  function checkString(props: { lens: Lens<string> }) {
-    return props;
-  }
-
-  checkString({ lens: email.cast<string>() });
-});