Add Pyrefly support

frontend/src/components/app-config/user-config-form.tsx

@@ -228,6 +228,7 @@ export const UserConfigForm: React.FC = () => {
       pylsp: true,
       ty: true,
       basedpyright: true,
+      pyrefly: true,
     };
   }
 
@@ -672,6 +673,48 @@ export const UserConfigForm: React.FC = () => {
                   </div>
                 )}
               />
+              <FormField
+                control={form.control}
+                name="language_servers.pyrefly.enabled"
+                render={({ field }) => (
+                  <div className="flex flex-col gap-1">
+                    <FormItem className={formItemClasses}>
+                      <FormLabel>
+                        <Badge variant="defaultOutline" className="mr-2">
+                          Beta
+                        </Badge>
+                        Pyrefly (
+                        <ExternalLink href="https://github.com/facebook/pyrefly">
+                          docs
+                        </ExternalLink>
+                        )
+                      </FormLabel>
+                      <FormControl>
+                        <Checkbox
+                          data-testid="pyrefly-checkbox"
+                          checked={field.value}
+                          disabled={field.disabled}
+                          onCheckedChange={(checked) => {
+                            field.onChange(Boolean(checked));
+                          }}
+                        />
+                      </FormControl>
+                      <FormMessage />
+                      <IsOverridden
+                        userConfig={config}
+                        name="language_servers.pyrefly.enabled"
+                      />
+                    </FormItem>
+                    {field.value && !capabilities.pyrefly && (
+                      <Banner kind="danger">
+                        Pyrefly is not available in your current environment.
+                        Please install <Kbd className="inline">pyrefly</Kbd> in
+                        your environment.
+                      </Banner>
+                    )}
+                  </div>
+                )}
+              />
               <FormField
                 control={form.control}
                 name="language_servers.ty.enabled"

frontend/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));
       }

frontend/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();

frontend/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 {

frontend/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}`);

marimo/_config/config.py

@@ -460,6 +460,17 @@ class TyLanguageServerConfig(TypedDict, total=False):
     enabled: bool
 
 
+@dataclass
+class PyreflyLanguageServerConfig(TypedDict, total=False):
+    """
+    Configuration options for Pyrefly Language Server.
+
+    Pyrefly handles completion, hover, go-to-definition, and diagnostics.
+    """
+
+    enabled: bool
+
+
 @dataclass
 class LanguageServersConfig(TypedDict, total=False):
     """Configuration options for language servers.
@@ -472,6 +483,7 @@ class LanguageServersConfig(TypedDict, total=False):
     pylsp: PythonLanguageServerConfig
     basedpyright: BasedpyrightServerConfig
     ty: TyLanguageServerConfig
+    pyrefly: PyreflyLanguageServerConfig
 
 
 @dataclass

marimo/_dependencies/dependencies.py

@@ -234,6 +234,7 @@ class DependencyManager:
     pylsp = Dependency("pylsp")
     basedpyright = Dependency("basedpyright")
     ty = Dependency("ty")
+    pyrefly = Dependency("pyrefly")
     pytest = Dependency("pytest")
     vegafusion = Dependency("vegafusion")
     vl_convert_python = Dependency("vl_convert")

marimo/_messaging/notification.py

@@ -249,13 +249,15 @@ class KernelCapabilitiesNotification(msgspec.Struct):
     pylsp: bool = False
     ty: bool = False
     basedpyright: bool = False
+    pyrefly: bool = False
 
     def __post_init__(self) -> None:
         # Only available in mac/linux
         self.terminal = not is_windows() and not is_pyodide()
         self.pylsp = DependencyManager.pylsp.has()
         self.basedpyright = DependencyManager.basedpyright.has()
         self.ty = DependencyManager.ty.has()
+        self.pyrefly = DependencyManager.pyrefly.has()
 
 
 class KernelReadyNotification(Notification, tag="kernel-ready"):

marimo/_schemas/generated/notifications.yaml

@@ -768,6 +768,9 @@ components:
         ty:
           default: false
           type: boolean
+        pyrefly:
+          default: false
+          type: boolean
       required: []
       title: KernelCapabilitiesNotification
       type: object

marimo/_server/lsp.py

@@ -607,6 +607,51 @@ def missing_binary_alert(self) -> AlertNotification:
         )
 
 
+class PyreflyServer(BaseLspServer):
+    id = "pyrefly"
+
+    def __init__(self, port: int) -> None:
+        super().__init__(port)
+        self.log_file = _loggers.get_log_directory() / "pyrefly-lsp.log"
+
+    async def start(self) -> Optional[AlertNotification]:
+        # Pyrefly is not required, so we don't want to alert or fail if it is not installed
+        if not DependencyManager.pyrefly.has():
+            LOGGER.debug("Pyrefly is not installed. Skipping LSP server.")
+            return None
+        return await super().start()
+
+    def validate_requirements(self) -> Union[str, Literal[True]]:
+        if not DependencyManager.pyrefly.has():
+            return "Pyrefly is missing. Install it with `pip install pyrefly`."
+        if not DependencyManager.which("node"):
+            return "node.js binary is missing. Install node at https://nodejs.org/."
+        return True
+
+    def get_command(self) -> list[str]:
+        from pyrefly.__main__ import get_pyrefly_bin  # type: ignore
+
+        lsp_bin = marimo_package_path() / "_lsp" / "index.cjs"
+        pyrefly_command = f"pyrefly:{get_pyrefly_bin()}"
+        return [
+            "node",
+            str(lsp_bin),
+            "--port",
+            str(self.port),
+            "--lsp",
+            pyrefly_command,
+            "--log-file",
+            str(self.log_file),
+        ]
+
+    def missing_binary_alert(self) -> AlertNotification:
+        return AlertNotification(
+            title="Pyrefly: Connection Error",
+            description="<span><a class='hyperlink' href='https://github.com/facebook/pyrefly'>Install pyrefly</a> for type checking support.</span>",
+            variant="danger",
+        )
+
+
 class NoopLspServer(LspServer):
     port: int = 0
     id: str = "noop"
@@ -635,6 +680,7 @@ class CompositeLspServer(LspServer):
         "pylsp": PyLspServer,
         "basedpyright": BasedpyrightServer,
         "ty": TyServer,
+        "pyrefly": PyreflyServer,
         "copilot": CopilotLspServer,
     }
 

packages/lsp/index.ts

@@ -121,6 +121,8 @@ export function parseTypedCommand(typedCommand: string): string[] {
       return ["node", binaryPath, "--stdio"];
     case "basedpyright":
       return [binaryPath, "--stdio"];
+    case "pyrefly":
+      return [binaryPath, "lsp"];
     case "ty":
       return [binaryPath, "server"];
     default:

packages/openapi/api.yaml

@@ -2184,6 +2184,9 @@ components:
         pylsp:
           default: false
           type: boolean
+        pyrefly:
+          default: false
+          type: boolean
         terminal:
           default: false
           type: boolean
@@ -2493,12 +2496,14 @@ components:
       type: object
     LanguageServersConfig:
       description: "Configuration options for language servers.\n\n    **Keys.**\n\
-        \n    - `pylsp`: the pylsp config"
+        \n    - `pylsp`: the pylsp config\n    - `pyrefly`: the pyrefly config\n    - `ty`: the ty config\n    - `basedpyright`: the basedpyright config"
       properties:
         basedpyright:
           $ref: '#/components/schemas/BasedpyrightServerConfig'
         pylsp:
           $ref: '#/components/schemas/PythonLanguageServerConfig'
+        pyrefly:
+          $ref: '#/components/schemas/PyreflyLanguageServerConfig'
         ty:
           $ref: '#/components/schemas/TyLanguageServerConfig'
       required: []
@@ -3491,6 +3496,17 @@ components:
       - tableName
       title: PreviewSQLTableRequest
       type: object
+    PyreflyLanguageServerConfig:
+      description: 'Configuration options for Pyrefly Language Server.
+
+
+        Pyrefly handles completion, hover, go-to-definition, and diagnostics.'
+      properties:
+        enabled:
+          type: boolean
+      required: []
+      title: PyreflyLanguageServerConfig
+      type: object
     PythonLanguageServerConfig:
       description: 'Configuration options for Python Language Server.
 
@@ -6217,4 +6233,3 @@ paths:
       summary: Submit login form
       tags:
       - auth
-

packages/openapi/src/api.ts

@@ -4387,6 +4387,8 @@ export interface components {
       /** @default false */
       pylsp?: boolean;
       /** @default false */
+      pyrefly?: boolean;
+      /** @default false */
       terminal?: boolean;
       /** @default false */
       ty?: boolean;
@@ -4560,10 +4562,14 @@ export interface components {
      *         **Keys.**
      *
      *         - `pylsp`: the pylsp config
+     *         - `pyrefly`: the pyrefly config
+     *         - `ty`: the ty config
+     *         - `basedpyright`: the basedpyright config
      */
     LanguageServersConfig: {
       basedpyright?: components["schemas"]["BasedpyrightServerConfig"];
       pylsp?: components["schemas"]["PythonLanguageServerConfig"];
+      pyrefly?: components["schemas"]["PyreflyLanguageServerConfig"];
       ty?: components["schemas"]["TyLanguageServerConfig"];
     };
     /** LayoutConfig */
@@ -5127,6 +5133,15 @@ export interface components {
       schema: string;
       tableName: string;
     };
+    /**
+     * PyreflyLanguageServerConfig
+     * @description Configuration options for Pyrefly Language Server.
+     *
+     *     Pyrefly handles completion, hover, go-to-definition, and diagnostics.
+     */
+    PyreflyLanguageServerConfig: {
+      enabled?: boolean;
+    };
     /**
      * PythonLanguageServerConfig
      * @description Configuration options for Python Language Server.