From 320f3f78ca86f692147776e4b1fa6f599dce2b11 Mon Sep 17 00:00:00 2001
From: Sterling Camden <rusty.camden@gmail.com>
Date: Fri, 17 Feb 2023 12:31:22 -0800
Subject: [PATCH] enable custom layouts and inter-field conditional display

---
 src/__tests__/createSchemaForm.test.tsx |  52 +++++++++++
 src/createSchemaForm.tsx                | 109 +++++++++++++-----------
 2 files changed, 112 insertions(+), 49 deletions(-)

diff --git a/src/__tests__/createSchemaForm.test.tsx b/src/__tests__/createSchemaForm.test.tsx
index 0f3e362..5eaf5ff 100644
--- a/src/__tests__/createSchemaForm.test.tsx
+++ b/src/__tests__/createSchemaForm.test.tsx
@@ -69,6 +69,58 @@ describe("createSchemaForm", () => {
     expect(screen.queryByTestId(testIds.textFieldTwo)).toBeTruthy();
     expect(screen.queryByTestId(testIds.booleanField)).toBeTruthy();
   });
+  it("should render a text field and a boolean field based on the mapping and schema into slots in a custom form", () => {
+    const testSchema = z.object({
+      textField: z.string(),
+      textFieldTwo: z.string(),
+      booleanField: z.string(),
+      t: z.string(),
+      t2: z.string(),
+      t3: z.string(),
+      t4: z.string(),
+      t5: z.string(),
+    });
+
+    const extraTestIds = {
+      extra1 : 'extra-form-fun',
+      extra2 : 'extra-form-fun2'
+    }
+    render(
+      <TestForm
+        onSubmit={() => {}}
+        schema={testSchema}
+        props={{
+          textField: {
+            testId: testIds.textField,
+          },
+          textFieldTwo: {
+            testId: testIds.textFieldTwo,
+          },
+          booleanField: {
+            testId: testIds.booleanField,
+          },
+        }}
+      >
+        {({renderedFields : {textField, booleanField, ...restFields}}) => {
+          return <>
+            <div data-testid={extraTestIds.extra1}>
+              {textField}
+            </div>
+            <div data-testid={extraTestIds.extra2}>
+              {booleanField}
+            </div>
+            {Object.values(restFields)}
+          </>
+        }}
+      </TestForm>
+    );
+      
+    expect(screen.queryByTestId(testIds.textField)).toBeTruthy();
+    expect(screen.queryByTestId(testIds.textFieldTwo)).toBeTruthy();
+    expect(screen.queryByTestId(testIds.booleanField)).toBeTruthy();
+    expect(screen.queryByTestId(extraTestIds.extra1)).toBeTruthy();
+    expect(screen.queryByTestId(extraTestIds.extra2 )).toBeTruthy();
+  });
   it("should render a text field and a boolean field based on the mapping and schema, unwrapping refine calls", () => {
     const testSchema = z.object({
       textField: z.string(),
diff --git a/src/createSchemaForm.tsx b/src/createSchemaForm.tsx
index cba1c92..91a86b8 100644
--- a/src/createSchemaForm.tsx
+++ b/src/createSchemaForm.tsx
@@ -1,6 +1,7 @@
 import React, {
   ForwardRefExoticComponent,
   Fragment,
+  FunctionComponent,
   ReactNode,
   RefAttributes,
   useRef,
@@ -237,6 +238,7 @@ export function createTsForm<
     renderAfter,
     renderBefore,
     form,
+    children : CustomChildrenComponent,
   }: {
     /**
      * A Zod Schema - An input field will be rendered for each property in the schema, based on the mapping passed to `createTsForm`
@@ -286,6 +288,9 @@ export function createTsForm<
      * ```
      */
     form?: UseFormReturn<z.infer<SchemaType>>;
+    children? : FunctionComponent<{renderedFields : {
+      [key in keyof z.infer<UnwrapEffects<SchemaType>>] : ReactNode
+    }}>
   } & RequireKeysWithRequiredChildren<{
     /**
      * Props to pass to the individual form components. The keys of `props` will be the names of your form properties in the form schema, and they will
@@ -392,58 +397,64 @@ export function createTsForm<
       });
     }
     const submitFn = handleSubmit(_submit);
-    return (
-      <FormProvider {..._form}>
-        <ActualFormComponent {...formProps} onSubmit={submitFn}>
-          {renderBefore && renderBefore({ submit: submitFn })}
-          {Object.keys(shape).map((key) => {
-            const type = shape[key] as RTFSupportedZodTypes;
-            const Component = getComponentForZodType(type, componentMap);
-            if (!Component) {
-              throw new Error(
-                noMatchingSchemaErrorMessage(key, type._def.typeName)
-              );
-            }
-            const meta = getMetaInformationForZodType(type);
+    type SchemaKey = keyof z.infer<UnwrapEffects<SchemaType>>;
+    const renderedFields = Object.keys(shape).reduce((accum, key : SchemaKey) => {
+      // we know this is a string but TS thinks it can be number and symbol so just in case stringify
+      const stringKey = key.toString();
+      const type = shape[key] as RTFSupportedZodTypes;
+      const Component = getComponentForZodType(type, componentMap);
+      if (!Component) {
+        throw new Error(
+          noMatchingSchemaErrorMessage(stringKey, type._def.typeName)
+        );
+      }
+      const meta = getMetaInformationForZodType(type);
 
-            const fieldProps = props && props[key] ? (props[key] as any) : {};
+      const fieldProps = props && props[key] ? (props[key] as any) : {};
 
-            const { beforeElement, afterElement } = fieldProps;
+      const { beforeElement, afterElement } = fieldProps;
 
-            const mergedProps = {
-              ...(propsMap.name && { [propsMap.name]: key }),
-              ...(propsMap.control && { [propsMap.control]: control }),
-              ...(propsMap.enumValues && {
-                [propsMap.enumValues]: meta.enumValues,
-              }),
-              ...(propsMap.descriptionLabel && {
-                [propsMap.descriptionLabel]: meta.description?.label,
-              }),
-              ...(propsMap.descriptionPlaceholder && {
-                [propsMap.descriptionPlaceholder]: meta.description?.placeholder,
-              }),
-              ...fieldProps,
-            };
-            const ctxLabel = meta.description?.label;
-            const ctxPlaceholder = meta.description?.placeholder;
-            return (
-              <Fragment key={key}>
-                {beforeElement}
-                <FieldContextProvider
-                  control={control}
-                  name={key}
-                  label={ctxLabel}
-                  placeholder={ctxPlaceholder}
-                  enumValues={meta.enumValues as string[] | undefined}
-                  addToCoerceUndefined={addToCoerceUndefined}
-                  removeFromCoerceUndefined={removeFromCoerceUndefined}
-                >
-                  <Component key={key} {...mergedProps} />
-                </FieldContextProvider>
-                {afterElement}
-              </Fragment>
-            );
-          })}
+      const mergedProps = {
+        ...(propsMap.name && { [propsMap.name]: key }),
+        ...(propsMap.control && { [propsMap.control]: control }),
+        ...(propsMap.enumValues && {
+          [propsMap.enumValues]: meta.enumValues,
+        }),
+        ...(propsMap.descriptionLabel && {
+          [propsMap.descriptionLabel]: meta.description?.label,
+        }),
+        ...(propsMap.descriptionPlaceholder && {
+          [propsMap.descriptionPlaceholder]: meta.description?.placeholder,
+        }),
+        ...fieldProps,
+      };
+      const ctxLabel = meta.description?.label;
+      const ctxPlaceholder = meta.description?.placeholder;
+      accum[key] = (
+        <Fragment key={stringKey}>
+          {beforeElement}
+          <FieldContextProvider
+            control={control}
+            name={stringKey}
+            label={ctxLabel}
+            placeholder={ctxPlaceholder}
+            enumValues={meta.enumValues as string[] | undefined}
+            addToCoerceUndefined={addToCoerceUndefined}
+            removeFromCoerceUndefined={removeFromCoerceUndefined}
+          >
+            <Component key={key} {...mergedProps} />
+          </FieldContextProvider>
+          {afterElement}
+        </Fragment>
+      );
+      return accum;
+    }, {} as Record<SchemaKey, React.ReactNode>);
+    const renderedFieldNodes = Object.values(renderedFields);
+    return (
+      <FormProvider {..._form}>
+        <ActualFormComponent {...formProps} onSubmit={submitFn} >
+          {renderBefore && renderBefore({ submit: submitFn })}
+          {CustomChildrenComponent ? <CustomChildrenComponent renderedFields={renderedFields}></CustomChildrenComponent> : renderedFieldNodes}
           {renderAfter && renderAfter({ submit: submitFn })}
         </ActualFormComponent>
       </FormProvider>