From 680ae4715f63c52d1c273649c6a66c4a00c46c39 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Fri, 9 May 2025 21:40:26 +0100
Subject: [PATCH 1/4] Format signatures in LSP hover and completion
 documentation

Remove short-circuit that was used to take a snapshot
---
 marimo/_server/lsp.py | 20 ++++++++++++++++++++
 pyproject.toml        |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/marimo/_server/lsp.py b/marimo/_server/lsp.py
index 138845c0a7b..c2930f2eab6 100644
--- a/marimo/_server/lsp.py
+++ b/marimo/_server/lsp.py
@@ -14,6 +14,7 @@
 from marimo._messaging.ops import Alert
 from marimo._server.utils import find_free_port
 from marimo._tracer import server_tracer
+from marimo._utils.formatter import FormatError, ruff
 from marimo._utils.paths import marimo_package_path
 
 LOGGER = _loggers.marimo_logger()
@@ -295,3 +296,22 @@ def any_lsp_server_running(config: MarimoConfig) -> bool:
         for server in language_servers.values()
     )
     return (copilot_enabled is not False) or language_servers_enabled
+
+
+if DependencyManager.pylsp.has():
+    from pylsp import hookimpl
+
+    def format_signature(signature: str) -> str:
+        try:
+            signature_as_func = f"def {signature.strip()}:\n    pass"
+            reformatted = ruff({"": signature_as_func}, "format")[""]
+            signature = reformatted.removeprefix("def ").removesuffix(
+                ":\n    pass"
+            )
+        except (ModuleNotFoundError, FormatError):
+            pass
+        return "```python\n" + signature + "\n```\n"
+
+    @hookimpl
+    def pylsp_signatures_to_markdown(signatures: list[str]) -> str:
+        return format_signature("\n".join(signatures))
diff --git a/pyproject.toml b/pyproject.toml
index b93242814d4..fdccd7d0017 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -78,6 +78,9 @@ homepage = "https://github.com/marimo-team/marimo"
 [project.entry-points."docstring_to_markdown"]
 marimo_converter = "marimo._utils.docs:MarimoConverter"
 
+[project.entry-points."pylsp"]
+marimo_plugin = "marimo._server.lsp"
+
 [project.optional-dependencies]
 sql = [
     "duckdb>=1.0.0",

From cb6e432f808cb9401b66a0cd2bd308f9a6b51542 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Sat, 10 May 2025 08:40:50 +0100
Subject: [PATCH 2/4] Fix typing issue

---
 pyproject.toml | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index fdccd7d0017..65890db5f7c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -195,6 +195,7 @@ dependencies = [
     "sqlglot>=23.4",
     "sqlalchemy>=2.0.40",
     "loro>=1.5.0",
+    "python-lsp-server>=1.12.1",
     "pandas-stubs>=1.5.3.230321",
     "pyiceberg>=0.9.0",
     "types-Pillow~=10.2.0.20240520",
@@ -473,6 +474,11 @@ exclude = [
 ]
 warn_unused_ignores = false
 
+[[tool.mypy.overrides]]
+module = "pylsp"
+# until https://github.com/python-lsp/python-lsp-server/pull/641 is released
+follow_untyped_imports = true
+
 [tool.pytest.ini_options]
 minversion = "6.0"
 addopts = "-ra -q -v --ignore tests/_cli/ipynb_data --ignore tests/_ast/codegen_data --ignore tests/_ast/app_data"

From 0606485476104690136fff3e6607d5633708b040 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Sat, 10 May 2025 09:38:01 +0100
Subject: [PATCH 3/4] Fix typing of cell ID

---
 marimo/_server/lsp.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/marimo/_server/lsp.py b/marimo/_server/lsp.py
index c2930f2eab6..cb14b41da24 100644
--- a/marimo/_server/lsp.py
+++ b/marimo/_server/lsp.py
@@ -14,6 +14,7 @@
 from marimo._messaging.ops import Alert
 from marimo._server.utils import find_free_port
 from marimo._tracer import server_tracer
+from marimo._types.ids import CellId_t
 from marimo._utils.formatter import FormatError, ruff
 from marimo._utils.paths import marimo_package_path
 
@@ -304,7 +305,10 @@ def any_lsp_server_running(config: MarimoConfig) -> bool:
     def format_signature(signature: str) -> str:
         try:
             signature_as_func = f"def {signature.strip()}:\n    pass"
-            reformatted = ruff({"": signature_as_func}, "format")[""]
+            dummy_cell_id = cast(CellId_t, "")
+            reformatted = ruff({dummy_cell_id: signature_as_func}, "format")[
+                dummy_cell_id
+            ]
             signature = reformatted.removeprefix("def ").removesuffix(
                 ":\n    pass"
             )

From cdf2df7c913c1160acd478fc151d777950e19ab6 Mon Sep 17 00:00:00 2001
From: krassowski <5832902+krassowski@users.noreply.github.com>
Date: Mon, 12 May 2025 09:52:12 +0100
Subject: [PATCH 4/4] Use the default formatter which will use ruff or black
 whichever available

---
 marimo/_server/lsp.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/marimo/_server/lsp.py b/marimo/_server/lsp.py
index cb14b41da24..c241e759836 100644
--- a/marimo/_server/lsp.py
+++ b/marimo/_server/lsp.py
@@ -15,7 +15,7 @@
 from marimo._server.utils import find_free_port
 from marimo._tracer import server_tracer
 from marimo._types.ids import CellId_t
-from marimo._utils.formatter import FormatError, ruff
+from marimo._utils.formatter import DefaultFormatter, FormatError
 from marimo._utils.paths import marimo_package_path
 
 LOGGER = _loggers.marimo_logger()
@@ -302,11 +302,13 @@ def any_lsp_server_running(config: MarimoConfig) -> bool:
 if DependencyManager.pylsp.has():
     from pylsp import hookimpl
 
+    formatter = DefaultFormatter(line_length=88)
+
     def format_signature(signature: str) -> str:
         try:
             signature_as_func = f"def {signature.strip()}:\n    pass"
             dummy_cell_id = cast(CellId_t, "")
-            reformatted = ruff({dummy_cell_id: signature_as_func}, "format")[
+            reformatted = formatter.format({dummy_cell_id: signature_as_func})[
                 dummy_cell_id
             ]
             signature = reformatted.removeprefix("def ").removesuffix(