diff --git a/engine/baml-schema-wasm/src/lib.rs b/engine/baml-schema-wasm/src/lib.rs index 243f94848..b9fb05baa 100644 --- a/engine/baml-schema-wasm/src/lib.rs +++ b/engine/baml-schema-wasm/src/lib.rs @@ -1,6 +1,7 @@ #[cfg(target_arch = "wasm32")] pub mod runtime_wasm; +use internal_baml_core::internal_baml_schema_ast::{format_schema, FormatOptions}; use std::env; use wasm_bindgen::prelude::*; @@ -9,3 +10,24 @@ pub fn version() -> String { // register_panic_hook(); env!("CARGO_PKG_VERSION").to_string() } + +#[wasm_bindgen] +pub fn format_document(path: String, text: String) -> Option { + log::info!("Trying to format document (rust): {}", path); + match format_schema( + &text, + FormatOptions { + indent_width: 2, + fail_on_unhandled_rule: false, + }, + ) { + Ok(formatted) => { + log::info!("Formatted document: {}", formatted); + Some(formatted) + } + Err(e) => { + log::error!("Failed to format document: {} {:?}", path, e); + None + } + } +} diff --git a/engine/cli/src/commands.rs b/engine/cli/src/commands.rs index 874ecb83e..273dea571 100644 --- a/engine/cli/src/commands.rs +++ b/engine/cli/src/commands.rs @@ -68,7 +68,12 @@ impl RuntimeCli { args.from = BamlRuntime::parse_baml_src_path(&args.from)?; t.block_on(async { args.run_async().await }) } - Commands::Format(args) => args.run(), + Commands::Format(args) => { + // We deliberately don't apply parse_baml_src_path here + // see format.rs for more details + // args.from = BamlRuntime::parse_baml_src_path(&args.from)?; + args.run() + } } } } diff --git a/engine/cli/src/format.rs b/engine/cli/src/format.rs index fd077d3eb..974604c24 100644 --- a/engine/cli/src/format.rs +++ b/engine/cli/src/format.rs @@ -1,6 +1,7 @@ use std::{fs, path::PathBuf}; use anyhow::Result; +use baml_runtime::{baml_src_files, BamlRuntime}; use clap::Args; use internal_baml_core::internal_baml_schema_ast::{format_schema, FormatOptions}; @@ -8,24 +9,56 @@ use internal_baml_core::internal_baml_schema_ast::{format_schema, FormatOptions} pub struct FormatArgs { #[arg(long, help = "path/to/baml_src", default_value = "./baml_src")] pub from: PathBuf, + + #[arg( + help = "Specific files to format. If none provided, formats all files in the baml_src directory" + )] + pub paths: Vec, + + #[arg( + short = 'n', + long = "dry-run", + help = "Write formatter changes to stdout instead of files", + default_value = "false" + )] + pub dry_run: bool, } impl FormatArgs { pub fn run(&self) -> Result<()> { - let source = fs::read_to_string(&self.from)?; - let formatted = format_schema( - &source, - FormatOptions { - indent_width: 4, - fail_on_unhandled_rule: false, - }, - )?; - - let mut to = self.from.clone(); - to.set_extension("formatted.baml"); - fs::write(&to, formatted)?; - - log::info!("Formatted {} to {}", self.from.display(), to.display()); + let paths = if self.paths.is_empty() { + // Usually this is done in commands.rs, but fmt is a special case + // because it doesn't need to actually load the BAML runtime to parse + // BAML files. + let from = BamlRuntime::parse_baml_src_path(&self.from)?; + baml_src_files(&from)? + } else { + self.paths.clone() + }; + + for path in paths.iter() { + let source = fs::read_to_string(&path)?; + match format_schema( + &source, + FormatOptions { + indent_width: 2, + fail_on_unhandled_rule: false, + }, + ) { + Ok(formatted) => { + if self.dry_run { + println!("{}", formatted); + } else { + fs::write(&path, formatted)?; + } + } + Err(e) => { + log::error!("Failed to format {}: {}", path.display(), e); + } + } + } + + log::info!("Formatted {} files", paths.len()); Ok(()) } diff --git a/tools/build b/tools/build index 0bac5abf9..2402b8781 100755 --- a/tools/build +++ b/tools/build @@ -137,7 +137,7 @@ case "$_path" in if [ "$_watch_mode" -eq 1 ]; then npx nodemon \ - --ext rs,hb,hbs,j2,toml,baml \ + --ext rs,hb,hbs,j2,toml \ --watch "${_repo_root}/engine" \ --ignore 'target' \ --exec "${command}" @@ -157,7 +157,7 @@ case "$_path" in if [ "$_watch_mode" -eq 1 ]; then npx nodemon \ - --ext rs,hb,hbs,j2,toml,baml \ + --ext rs,hb,hbs,j2,toml \ --watch "${_repo_root}/engine" \ --ignore 'target/**' \ --exec "${command}" diff --git a/typescript/vscode-ext/packages/language-server/src/server.ts b/typescript/vscode-ext/packages/language-server/src/server.ts index 1520d0057..771318138 100644 --- a/typescript/vscode-ext/packages/language-server/src/server.ts +++ b/typescript/vscode-ext/packages/language-server/src/server.ts @@ -140,7 +140,7 @@ export function startServer(options?: LSOptions): void { capabilities: { textDocumentSync: TextDocumentSyncKind.Full, definitionProvider: true, - documentFormattingProvider: false, + documentFormattingProvider: true, completionProvider: { resolveProvider: false, triggerCharacters: ['@', '"', '.'], @@ -527,6 +527,21 @@ export function startServer(options?: LSOptions): void { } }) + connection.onDocumentFormatting((params: DocumentFormattingParams) => { + try { + const doc = getDocument(params.textDocument.uri) + if (doc) { + const formatted = BamlWasm.format_document(doc.uri, doc.getText()) + if (formatted) { + return [TextEdit.replace(Range.create(doc.positionAt(0), doc.positionAt(doc.getText().length)), formatted)] + } + } + return [] + } catch (e) { + console.error(`Error occurred while formatting document:\n${e}`) + } + }) + connection.onCodeLens((params: CodeLensParams) => { try { const document = getDocument(params.textDocument.uri) diff --git a/typescript/vscode-ext/packages/package.json b/typescript/vscode-ext/packages/package.json index 676e5f684..4da434aa2 100644 --- a/typescript/vscode-ext/packages/package.json +++ b/typescript/vscode-ext/packages/package.json @@ -154,6 +154,11 @@ "command": "baml.selectTestCase", "title": "Select Test Case", "category": "Baml" + }, + { + "command": "baml.setDefaultFormatter", + "title": "Set Default Formatter", + "category": "Baml" } ] }, diff --git a/typescript/vscode-ext/packages/vscode/src/extension.ts b/typescript/vscode-ext/packages/vscode/src/extension.ts index 7809e20bf..345b01d16 100644 --- a/typescript/vscode-ext/packages/vscode/src/extension.ts +++ b/typescript/vscode-ext/packages/vscode/src/extension.ts @@ -354,6 +354,12 @@ export function activate(context: vscode.ExtensionContext) { } }) + const config = vscode.workspace.getConfiguration('editor', { languageId: 'baml' }) + if (!config.get('defaultFormatter')) { + // TODO: once the BAML formatter is stable, we should auto-prompt people to set it as the default formatter. + // void vscode.commands.executeCommand('baml.setDefaultFormatter') + } + // Listen for messages from the webview plugins.map(async (plugin) => { @@ -394,7 +400,6 @@ export function deactivate(): void { } server?.close() } - class DiagnosticCodeActionProvider implements vscode.CodeActionProvider { public provideCodeActions( document: vscode.TextDocument, diff --git a/typescript/vscode-ext/packages/vscode/src/plugins/language-server/index.ts b/typescript/vscode-ext/packages/vscode/src/plugins/language-server/index.ts index 5ab3a973b..88e40bb58 100644 --- a/typescript/vscode-ext/packages/vscode/src/plugins/language-server/index.ts +++ b/typescript/vscode-ext/packages/vscode/src/plugins/language-server/index.ts @@ -382,6 +382,49 @@ const plugin: BamlVSCodePlugin = { } }, ), + + commands.registerCommand('baml.setDefaultFormatter', async () => { + enum AutoFormatChoice { + Yes = 'Yes (always)', + OnlyInWorkspace = 'Yes (in workspace)', + No = 'No', + } + const selection = await vscode.window.showInformationMessage( + 'Would you like to auto-format BAML files on save?', + AutoFormatChoice.Yes, + AutoFormatChoice.OnlyInWorkspace, + AutoFormatChoice.No, + ) + if (selection === AutoFormatChoice.No) { + return + } + + const config = vscode.workspace.getConfiguration('editor', { languageId: 'baml' }) + + const configTarget = + selection === AutoFormatChoice.Yes ? vscode.ConfigurationTarget.Global : vscode.ConfigurationTarget.Workspace + const overrideInLanguage = true + + for (const [key, value] of Object.entries({ + defaultFormatter: 'Boundary.baml-extension', + formatOnSave: true, + })) { + await config.update(key, value, configTarget, overrideInLanguage) + } + + switch (selection) { + case AutoFormatChoice.Yes: + vscode.window.showInformationMessage( + 'BAML files will now be auto-formatted on save (updated user settings).', + ) + break + case AutoFormatChoice.OnlyInWorkspace: + vscode.window.showInformationMessage( + 'BAML files will now be auto-formatted on save (updated workspace settings).', + ) + break + } + }), ) activateClient(context, serverOptions, clientOptions)