@marimo-team/frontend 0.19.11-dev2 → 0.19.11-dev4

package/src/core/codemirror/language/languages/python.ts

@@ -175,6 +175,44 @@ const tyLspClient = once((_: LSPConfig) => {
   return notebookClient;
 });
 
+const pyreflyClient = once(
+  (lspConfig: LSPConfig & { diagnostics: DiagnosticsConfig }) => {
+    let resyncCallback: (() => Promise<void>) | undefined;
+
+    const transport = createTransport("pyrefly", async () => {
+      await resyncCallback?.();
+    });
+
+    const lspClientOpts = {
+      transport,
+      rootUri: getLSPDocumentRootUri(),
+      workspaceFolders: [],
+    };
+
+    // We wrap the client in a NotebookLanguageServerClient to add some
+    // additional functionality to handle multiple cells
+    const notebookClient = new NotebookLanguageServerClient(
+      new LanguageServerClient({
+        ...lspClientOpts,
+        initializationOptions: {
+          pyrefly: {
+            displayTypeErrors:
+              (lspConfig.diagnostics?.enabled ?? false)
+                ? "force-on"
+                : "force-off",
+          },
+        },
+      }),
+      {},
+    );
+
+    // Set the resync callback now that the client exists
+    resyncCallback = () => notebookClient.resyncAllDocuments();
+
+    return notebookClient;
+  },
+);
+
 const pyrightClient = once((_: LSPConfig) => {
   let resyncCallback: (() => Promise<void>) | undefined;
 
@@ -263,6 +301,9 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
       if (lspConfig?.ty?.enabled && hasCapability("ty")) {
         clients.push(tyLspClient(lspConfig));
       }
+      if (lspConfig?.pyrefly?.enabled && hasCapability("pyrefly")) {
+        clients.push(pyreflyClient(lspConfig));
+      }
       if (lspConfig?.basedpyright?.enabled && hasCapability("basedpyright")) {
         clients.push(pyrightClient(lspConfig));
       }

package/src/core/codemirror/lsp/transports.ts

@@ -14,7 +14,7 @@ import { getRuntimeManager } from "../../runtime/config";
  * @returns The transport.
  */
 export function createTransport(
-  serverName: "pylsp" | "basedpyright" | "copilot" | "ty",
+  serverName: "pylsp" | "basedpyright" | "copilot" | "ty" | "pyrefly",
   onReconnect?: () => Promise<void>,
 ) {
   const runtimeManager = getRuntimeManager();

package/src/core/config/capabilities.ts

@@ -8,6 +8,7 @@ export const capabilitiesAtom = atom<Capabilities>({
   pylsp: false,
   basedpyright: false,
   ty: false,
+  pyrefly: false,
 });
 
 export function hasCapability(key: keyof Capabilities): boolean {

package/src/core/runtime/runtime.ts

@@ -140,7 +140,7 @@ export class RuntimeManager {
   /**
    * The URL of the copilot server.
    */
-  getLSPURL(lsp: "pylsp" | "basedpyright" | "copilot" | "ty"): URL {
+  getLSPURL(lsp: "pylsp" | "basedpyright" | "copilot" | "ty" | "pyrefly"): URL {
     if (lsp === "copilot") {
       // For copilot, don't include any query parameters
       const url = this.formatWsURL(`/lsp/${lsp}`);