Fix printing when block attributes are mixed in with fields

MrLeebo · Mar 9, 2024 · 6585830 · 6585830
1 parent 147589a
commit 6585830
Show file tree

Hide file tree

Showing 10 changed files with 1,024 additions and 22 deletions.
diff --git a/src/parser.ts b/src/parser.ts
@@ -95,7 +95,7 @@ export class PrismaParser extends CstParser {
       ]);
     });
     this.MANY(() => {
-      this.SUBRULE(this.attribute, { LABEL: 'attributeList' });
+      this.SUBRULE(this.fieldAttribute, { LABEL: 'attributeList' });
     });
     this.OPTION2(() => {
       this.CONSUME(lexer.Comment, { LABEL: 'comment' });
@@ -125,7 +125,7 @@ export class PrismaParser extends CstParser {
             GATE: () => isObject,
             ALT: () => this.SUBRULE(this.property, { LABEL: 'list' }),
           },
-          { ALT: () => this.SUBRULE(this.attribute, { LABEL: 'list' }) },
+          { ALT: () => this.SUBRULE(this.blockAttribute, { LABEL: 'list' }) },
           {
             GATE: () => isObject,
             ALT: () => this.SUBRULE(this.field, { LABEL: 'list' }),
@@ -152,18 +152,10 @@ export class PrismaParser extends CstParser {
       this.CONSUME(lexer.Comment, { LABEL: 'comment' });
     });
   });
-  private attribute = this.RULE('attribute', () => {
-    this.OR1([
-      {
-        ALT: () =>
-          this.CONSUME(lexer.BlockAttribute, { LABEL: 'blockAttribute' }),
-      },
-      {
-        ALT: () =>
-          this.CONSUME(lexer.FieldAttribute, { LABEL: 'fieldAttribute' }),
-      },
-    ]);
-    this.OR2([
+
+  private fieldAttribute = this.RULE('fieldAttribute', () => {
+    this.CONSUME(lexer.FieldAttribute, { LABEL: 'fieldAttribute' });
+    this.OR([
       {
         ALT: () => {
           this.CONSUME1(lexer.Identifier, { LABEL: 'groupName' });
@@ -188,6 +180,33 @@ export class PrismaParser extends CstParser {
     });
   });
 
+  private blockAttribute = this.RULE('blockAttribute', () => {
+    this.CONSUME(lexer.BlockAttribute, { LABEL: 'blockAttribute' }),
+      this.OR([
+        {
+          ALT: () => {
+            this.CONSUME1(lexer.Identifier, { LABEL: 'groupName' });
+            this.CONSUME(lexer.Dot);
+            this.CONSUME2(lexer.Identifier, { LABEL: 'attributeName' });
+          },
+        },
+        {
+          ALT: () => this.CONSUME(lexer.Identifier, { LABEL: 'attributeName' }),
+        },
+      ]);
+
+    this.OPTION(() => {
+      this.CONSUME(lexer.LRound);
+      this.MANY_SEP({
+        SEP: lexer.Comma,
+        DEF: () => {
+          this.SUBRULE(this.attributeArg);
+        },
+      });
+      this.CONSUME(lexer.RRound);
+    });
+  });
+
   private attributeArg = this.RULE('attributeArg', () => {
     this.OR([
       {

diff --git a/src/printSchema.ts b/src/printSchema.ts
@@ -115,7 +115,47 @@ generator ${generator.name} {
 }
 
 function printObject(object: Types.Object) {
-  const children = computePropertyFormatting(object.properties);
+  const props = [...object.properties];
+
+  // If block attributes are declared in the middle of the block, move them to
+  // the bottom of the list.
+  let blockAttributeMoved = false;
+  props.sort((a, b) => {
+    if (
+      a.type === 'attribute' &&
+      a.kind === 'object' &&
+      (b.type !== 'attribute' ||
+        (b.type === 'attribute' && b.kind !== 'object'))
+    ) {
+      blockAttributeMoved = true;
+      return 1;
+    }
+
+    if (
+      b.type === 'attribute' &&
+      b.kind === 'object' &&
+      (a.type !== 'attribute' ||
+        (a.type === 'attribute' && a.kind !== 'object'))
+    ) {
+      blockAttributeMoved = true;
+      return -1;
+    }
+
+    return 0;
+  });
+
+  // Insert a break between the block attributes and the file if the block
+  // attributes are too close to the model's fields
+  const attrIndex = props.findIndex(
+    (item) => item.type === 'attribute' && item.kind === 'object'
+  );
+
+  const needsSpace = !['break', 'comment'].includes(props[attrIndex - 1]?.type);
+  if (blockAttributeMoved && needsSpace) {
+    props.splice(attrIndex, 0, { type: 'break' });
+  }
+
+  const children = computePropertyFormatting(props);
 
   return `
 ${object.type} ${object.name} {

diff --git a/src/visitor.ts b/src/visitor.ts
@@ -92,7 +92,7 @@ export const VisitorClassFactory = (
       return this.maybeAppendLocationData(data, comment);
     }
 
-    block(ctx: CstNode & { list: CstNode[] }): Types.Block[] {
+    block(ctx: CstNode & { list: CstNode[] }): BlockList {
       return ctx.list?.map((item) => this.visit([item]));
     }
 
@@ -139,9 +139,8 @@ export const VisitorClassFactory = (
       );
     }
 
-    attribute(
+    fieldAttribute(
       ctx: CstNode & {
-        blockAttribute: [IToken];
         fieldAttribute: [IToken];
         groupName: [IToken];
         attributeName: [IToken];
@@ -152,16 +151,47 @@ export const VisitorClassFactory = (
       const [group] = ctx.groupName || [{}];
       const args =
         ctx.attributeArg && ctx.attributeArg.map((attr) => this.visit(attr));
-      const kind = ctx.blockAttribute != null ? 'object' : 'field';
       const data = {
         type: 'attribute',
         name: name.image,
-        kind,
+        kind: 'field',
+        group: group.image,
+        args,
+      } as const;
+      return this.maybeAppendLocationData(
+        data,
+        name,
+        ...ctx.fieldAttribute,
+        group
+      );
+    }
+
+    blockAttribute(
+      ctx: CstNode & {
+        blockAttribute: [IToken];
+        groupName: [IToken];
+        attributeName: [IToken];
+        attributeArg: CstNode[];
+      }
+    ): Types.Attr | null {
+      const [name] = ctx.attributeName;
+      const [group] = ctx.groupName || [{}];
+      const args =
+        ctx.attributeArg && ctx.attributeArg.map((attr) => this.visit(attr));
+      const data = {
+        type: 'attribute',
+        name: name.image,
+        kind: 'object',
         group: group.image,
         args,
       } as const;
-      const attrs = kind === 'object' ? ctx.blockAttribute : ctx.fieldAttribute;
-      return this.maybeAppendLocationData(data, name, ...attrs, group);
+
+      return this.maybeAppendLocationData(
+        data,
+        name,
+        ...ctx.blockAttribute,
+        group
+      );
     }
 
     attributeArg(ctx: CstNode & { value: [CstNode] }): Types.AttributeArgument {
@@ -233,5 +263,14 @@ export const VisitorClassFactory = (
   };
 };
 
+type BlockList = Array<
+  | Types.Comment
+  | Types.Property
+  | Types.Attribute
+  | Types.Field
+  | Types.Enum
+  | Types.Assignment
+  | Types.Break
+>;
 export const DefaultVisitorClass = VisitorClassFactory(defaultParser);
 export const defaultVisitor = new DefaultVisitorClass();
diff --git a/test/PrismaSchemaBuilder.test.ts b/test/PrismaSchemaBuilder.test.ts
@@ -451,6 +451,7 @@ describe('PrismaSchemaBuilder', () => {
       model User {
         firstName String
         lastName  String
+
         @@id([firstName, lastName])
       }
       "
@@ -469,6 +470,7 @@ describe('PrismaSchemaBuilder', () => {
       model Project {
         code   String
         client String
+
         @@unique([code, client])
       }
       "