feat: support prisma format indentation (#1256)

Co-authored-by: ymc9 <[email protected]>
zenstackhq · Apr 15, 2024 · b0f5d3b · b0f5d3b
1 parent c37bf92
commit b0f5d3b
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 7 deletions.
diff --git a/packages/schema/src/cli/actions/format.ts b/packages/schema/src/cli/actions/format.ts
@@ -6,7 +6,12 @@ import ora from 'ora';
 import { CliError } from '../cli-error';
 import { formatDocument, getDefaultSchemaLocation } from '../cli-util';
 
-export async function format(_projectPath: string, options: { schema: string }) {
+type Options = {
+    schema: string;
+    prismaStyle?: boolean;
+};
+
+export async function format(_projectPath: string, options: Options) {
     const version = getVersion();
     console.log(colors.bold(`⌛️ ZenStack CLI v${version}`));
 
@@ -18,7 +23,7 @@ export async function format(_projectPath: string, options: { schema: string })
 
     const spinner = ora(`Formatting ${schemaFile}`).start();
     try {
-        const formattedDoc = await formatDocument(schemaFile);
+        const formattedDoc = await formatDocument(schemaFile, options.prismaStyle);
         await writeFile(schemaFile, formattedDoc);
         spinner.succeed();
     } catch (e) {

diff --git a/packages/schema/src/cli/cli-util.ts b/packages/schema/src/cli/cli-util.ts
@@ -266,7 +266,7 @@ export async function checkNewVersion() {
     }
 }
 
-export async function formatDocument(fileName: string) {
+export async function formatDocument(fileName: string, isPrismaStyle = true) {
     const services = createZModelServices(NodeFileSystem).ZModel;
     const extensions = services.LanguageMetaData.fileExtensions;
     if (!extensions.includes(path.extname(fileName))) {
@@ -279,6 +279,8 @@ export async function formatDocument(fileName: string) {
 
     const formatter = services.lsp.Formatter as ZModelFormatter;
 
+    formatter.setPrismaStyle(isPrismaStyle);
+
     const identifier = { uri: document.uri.toString() };
     const options = formatter.getFormatOptions() ?? {
         insertSpaces: true,

diff --git a/packages/schema/src/cli/index.ts b/packages/schema/src/cli/index.ts
@@ -128,6 +128,7 @@ export function createProgram() {
         .command('format')
         .description('Format a ZenStack schema file.')
         .addOption(schemaOption)
+        .option('--no-prisma-style', 'do not use prisma style')
         .action(formatAction);
 
     // make sure config is loaded before actions run

diff --git a/packages/schema/src/language-server/zmodel-formatter.ts b/packages/schema/src/language-server/zmodel-formatter.ts
@@ -5,12 +5,31 @@ import { FormattingOptions, Range, TextEdit } from 'vscode-languageserver';
 
 export class ZModelFormatter extends AbstractFormatter {
     private formatOptions?: FormattingOptions;
+    private isPrismaStyle = true;
     protected format(node: AstNode): void {
         const formatter = this.getNodeFormatter(node);
+
         if (ast.isDataModelField(node)) {
-            formatter.property('type').prepend(Formatting.oneSpace());
-            if (node.attributes.length > 0) {
-                formatter.properties('attributes').prepend(Formatting.oneSpace());
+            if (this.isPrismaStyle && ast.isDataModel(node.$container)) {
+                const dataModel = node.$container;
+
+                const compareFn = (a: number, b: number) => b - a;
+                const maxNameLength = dataModel.fields.map((x) => x.name.length).sort(compareFn)[0];
+                const maxTypeLength = dataModel.fields.map(this.getFieldTypeLength).sort(compareFn)[0];
+
+                formatter.property('type').prepend(Formatting.spaces(maxNameLength - node.name.length + 1));
+                if (node.attributes.length > 0) {
+                    formatter
+                        .node(node.attributes[0])
+                        .prepend(Formatting.spaces(maxTypeLength - this.getFieldTypeLength(node) + 1));
+
+                    formatter.nodes(...node.attributes.slice(1)).prepend(Formatting.oneSpace());
+                }
+            } else {
+                formatter.property('type').prepend(Formatting.oneSpace());
+                if (node.attributes.length > 0) {
+                    formatter.properties('attributes').prepend(Formatting.oneSpace());
+                }
             }
         } else if (ast.isDataModelFieldAttribute(node)) {
             formatter.keyword('(').surround(Formatting.noSpace());
@@ -52,4 +71,22 @@ export class ZModelFormatter extends AbstractFormatter {
     public getIndent() {
         return 1;
     }
+
+    public setPrismaStyle(isPrismaStyle: boolean) {
+        this.isPrismaStyle = isPrismaStyle;
+    }
+
+    private getFieldTypeLength(field: ast.DataModelField) {
+        let length = (field.type.type || field.type.reference?.$refText)!.length;
+
+        if (field.type.optional) {
+            length += 1;
+        }
+
+        if (field.type.array) {
+            length += 2;
+        }
+
+        return length;
+    }
 }
diff --git a/packages/schema/tests/schema/formatter.test.ts b/packages/schema/tests/schema/formatter.test.ts
@@ -25,7 +25,7 @@ plugin swrHooks {
     output = 'lib/hooks'
 }
 model User {
-    id String @id
+    id   String  @id
     name String?
 }
 enum Role {

diff --git a/tests/integration/tests/cli/format.test.ts b/tests/integration/tests/cli/format.test.ts
@@ -35,6 +35,32 @@ generator client {
 model Post {
     id Int @id() @default(autoincrement())
     users User[]
+}`;
+        // set up schema
+        fs.writeFileSync('schema.zmodel', model, 'utf-8');
+        const program = createProgram();
+        await program.parseAsync(['format', '--no-prisma-style'], { from: 'user' });
+
+        expect(fs.readFileSync('schema.zmodel', 'utf-8')).toEqual(formattedModel);
+    });
+
+    it('prisma format', async () => {
+        const model = `
+        datasource db {provider="sqlite" url="file:./dev.db"}
+        generator client {provider = "prisma-client-js"}
+        model Post {id Int @id() @default(autoincrement())users User[]}`;
+
+        const formattedModel = `
+datasource db {
+    provider="sqlite"
+    url="file:./dev.db"
+}
+generator client {
+    provider = "prisma-client-js"
+}
+model Post {
+    id    Int    @id() @default(autoincrement())
+    users User[]
 }`;
         // set up schema
         fs.writeFileSync('schema.zmodel', model, 'utf-8');