✨ feat: better optional handling

m1212e · Mar 25, 2024 · 7acb6f3 · 7acb6f3
1 parent 09022f9
commit 7acb6f3
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 20 deletions.
diff --git a/prisma/schema.prisma b/prisma/schema.prisma
@@ -8,7 +8,7 @@ generator client {
 }
 
 generator prismabox {
-  provider = "node ./dist/cli.js"
+  provider  = "node ./dist/cli.js"
   dataModel = true
 }
 
@@ -33,11 +33,16 @@ model User {
 model Post {
   id        Int      @id @default(autoincrement())
   createdAt DateTime @default(now())
-  User   User? @relation(fields: [userId], references: [id])
+  User      User?    @relation(fields: [userId], references: [id])
   /// @prismabox.options{max: 10}
   /// this is the user id
-  userId Int?
-  status Status
+  userId    Int?
+
+  status  Status
+  status2 Status?
+
+  blah  String
+  blah2 String?
 
   successorId Int?  @unique
   successor   Post? @relation("history", fields: [successorId], references: [id])

diff --git a/src/generator/dataModel.ts b/src/generator/dataModel.ts
@@ -1,7 +1,11 @@
 import type { DMMF } from "@prisma/generator-helper";
 import { typeboxImportVariableName } from "./typeboxImport";
 import { Annotation, parseDocumentation } from "./documentation";
-import { PrimitiveField, isPrimitivePrismaFieldType } from "./plainModel";
+import {
+  NullableVariant,
+  PrimitiveField,
+  isPrimitivePrismaFieldType,
+} from "./plainModel";
 import type { Models } from "../util/modelMap";
 import { RelationField } from "./relationModel";
 
@@ -11,9 +15,23 @@ export function enableDataModel() {
 }
 
 export function DataModel(
+  data: Pick<DMMF.Model, "fields" | "documentation">,
+  referenceableEnums: Models
+) {
+  return internal(data, referenceableEnums, false);
+}
+
+export function DataModelOptional(
+  data: Pick<DMMF.Model, "fields" | "documentation">,
+  referenceableEnums: Models
+) {
+  return internal(data, referenceableEnums, true);
+}
+
+function internal(
   data: Pick<DMMF.Model, "fields" | "documentation">,
   referenceableEnums: Models,
-  optional = false
+  optional: boolean
 ) {
   if (!enabled) return undefined;
   const modelDoc = parseDocumentation(data.documentation);
@@ -43,7 +61,11 @@ export function DataModel(
           fieldType: field.type,
           list: field.isList,
           name: field.name,
-          optional,
+          optional: field.isRequired
+            ? optional
+              ? NullableVariant.OPTIONAL
+              : NullableVariant.REQUIRED
+            : NullableVariant.NULLABLE,
           options: doc.options,
           referenceableModels: referenceableEnums,
         });
@@ -77,7 +99,11 @@ export function DataModel(
         fieldType: field.type as any, // we checked earlier if it's a primitive type
         list: field.isList,
         name: field.name,
-        optional,
+        optional: field.isRequired
+            ? optional
+              ? NullableVariant.OPTIONAL
+              : NullableVariant.REQUIRED
+            : NullableVariant.NULLABLE,
         options: parseDocumentation(field.documentation).options,
       });
     })

diff --git a/src/generator/plainModel.ts b/src/generator/plainModel.ts
@@ -2,9 +2,15 @@ import type { DMMF } from "@prisma/generator-helper";
 import { typeboxImportVariableName } from "./typeboxImport";
 import { Annotation, parseDocumentation } from "./documentation";
 
+export enum NullableVariant {
+  NULLABLE = 0, // null, undefined or the value
+  OPTIONAL = 1, // undefined or the value
+  REQUIRED = 2, // the value
+}
+
 export function PlainModel(data: Pick<DMMF.Model, "fields" | "documentation">) {
   const modelDoc = parseDocumentation(data.documentation);
-  
+
   if (modelDoc.annotations.includes(Annotation.HIDDEN)) return undefined;
 
   const fields = data.fields
@@ -17,7 +23,9 @@ export function PlainModel(data: Pick<DMMF.Model, "fields" | "documentation">) {
         name: field.name,
         fieldType: field.type,
         list: field.isList,
-        optional: !field.isRequired,
+        optional: field.isRequired
+          ? NullableVariant.REQUIRED
+          : NullableVariant.NULLABLE,
         options: doc.options,
       });
     })
@@ -63,13 +71,15 @@ export function PrimitiveField({
   fieldType: PrimitivePrismaFieldType;
   options: string;
   name: string;
-  optional: boolean;
+  optional: NullableVariant;
   list: boolean;
 }) {
   let ret = `${name}: `;
 
-  if (optional) {
+  if (optional === NullableVariant.NULLABLE) {
     ret += "Nullable(";
+  } else if (optional === NullableVariant.OPTIONAL) {
+    ret += `${typeboxImportVariableName}.Optional(`;
   }
 
   if (list) {
@@ -90,7 +100,10 @@ export function PrimitiveField({
     ret += `${typeboxImportVariableName}.Boolean(${options})`;
   } else throw new Error("Invalid type for primitive generation");
 
-  if (optional) {
+  if (
+    optional === NullableVariant.NULLABLE ||
+    optional === NullableVariant.OPTIONAL
+  ) {
     ret += ")";
   }
 

diff --git a/src/generator/relationModel.ts b/src/generator/relationModel.ts
@@ -2,7 +2,7 @@ import type { DMMF } from "@prisma/generator-helper";
 import { typeboxImportVariableName } from "./typeboxImport";
 import { Annotation, parseDocumentation } from "./documentation";
 import type { Models } from "../util/modelMap";
-import { isPrimitivePrismaFieldType } from "./plainModel";
+import { NullableVariant, isPrimitivePrismaFieldType } from "./plainModel";
 
 export function RelationModel(
   data: Pick<DMMF.Model, "fields" | "documentation">,
@@ -21,7 +21,9 @@ export function RelationModel(
         name: field.name,
         fieldType: field.type,
         list: field.isList,
-        optional: !field.isRequired,
+        optional: field.isRequired
+          ? NullableVariant.REQUIRED
+          : NullableVariant.NULLABLE,
         options: doc.options,
         referenceableModels,
       });
@@ -42,7 +44,7 @@ export function RelationField({
   fieldType: string;
   options: string;
   name: string;
-  optional: boolean;
+  optional: NullableVariant;
   list: boolean;
   referenceableModels: Models;
 }) {
@@ -54,8 +56,10 @@ export function RelationField({
 
   let ret = `${name}: `;
 
-  if (optional) {
+  if (optional === NullableVariant.NULLABLE) {
     ret += "Nullable(";
+  } else if (optional === NullableVariant.OPTIONAL) {
+    ret += `${typeboxImportVariableName}.Optional(`;
   }
 
   if (list) {
@@ -71,7 +75,10 @@ export function RelationField({
   }
   ret += referencedFieldModel;
 
-  if (optional) {
+  if (
+    optional === NullableVariant.NULLABLE ||
+    optional === NullableVariant.OPTIONAL
+  ) {
     ret += ")";
   }
 

diff --git a/src/index.ts b/src/index.ts
@@ -15,7 +15,7 @@ import { setAdditionalProperties } from "./generator/documentation";
 import { Composite } from "./generator/composite";
 import { WhereModel } from "./generator/whereModel";
 import { Nullable } from "./generator/nullable";
-import { DataModel, enableDataModel } from "./generator/dataModel";
+import { DataModel, DataModelOptional, enableDataModel } from "./generator/dataModel";
 
 generatorHandler({
   onManifest() {
@@ -113,7 +113,7 @@ generatorHandler({
 
     optionalDataTasks.push(
       ...options.dmmf.datamodel.models.map(async (e) => {
-        const model = DataModel(e, enumTypes, true);
+        const model = DataModelOptional(e, enumTypes);
         if (model) {
           optionalDataTypes.set(e.name, model);
         }