Format signatures in LSP hover and completion documentation

marimo/_server/lsp.py

@@ -14,6 +14,8 @@
 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 DefaultFormatter, FormatError
 from marimo._utils.paths import marimo_package_path
 
 LOGGER = _loggers.marimo_logger()
@@ -295,3 +297,27 @@ 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
+
+    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 = formatter.format({dummy_cell_id: signature_as_func})[
+                dummy_cell_id
+            ]
+            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))

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",
@@ -192,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",
@@ -470,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"