Implement workspace search

Author: niilohlin

pylsp/config/schema.json

@@ -242,21 +242,26 @@
       "default": true,
       "description": "Enable or disable the plugin."
     },
-    "pylsp.plugins.jedi_symbols.enabled": {
+    "pylsp.plugins.jedi_document_symbols.enabled": {
       "type": "boolean",
       "default": true,
       "description": "Enable or disable the plugin."
     },
-    "pylsp.plugins.jedi_symbols.all_scopes": {
+    "pylsp.plugins.jedi_document_symbols.all_scopes": {
       "type": "boolean",
       "default": true,
       "description": "If True lists the names of all scopes instead of only the module namespace."
     },
-    "pylsp.plugins.jedi_symbols.include_import_symbols": {
+    "pylsp.plugins.jedi_document_symbols.include_import_symbols": {
       "type": "boolean",
       "default": true,
       "description": "If True includes symbols imported from other libraries."
     },
+    "pylsp.plugins.jedi_workspace_symbols.enabled": {
+      "type": "boolean",
+      "default": true,
+      "description": "If True includes workspace symbols."
+    },
     "pylsp.plugins.mccabe.enabled": {
       "type": "boolean",
       "default": true,

pylsp/hookspecs.py

@@ -68,6 +68,11 @@ def pylsp_execute_command(config, workspace, command, arguments):
     pass
 
 
+@hookspec
+def pylsp_workspace_symbol(config, workspace, document, query):
+    pass
+
+
 @hookspec
 def pylsp_experimental_capabilities(config, workspace):
     pass

pylsp/plugins/symbols.py renamed to pylsp/plugins/document_symbols.py

@@ -12,7 +12,7 @@
 
 @hookimpl
 def pylsp_document_symbols(config, document):
-    symbols_settings = config.plugin_settings("jedi_symbols")
+    symbols_settings = config.plugin_settings("jedi_document_symbols")
     all_scopes = symbols_settings.get("all_scopes", True)
     add_import_symbols = symbols_settings.get("include_import_symbols", True)
     definitions = document.jedi_names(all_scopes=all_scopes)

pylsp/plugins/workspace_symbol.py

@@ -0,0 +1,50 @@
+# Copyright 2017-2020 Palantir Technologies, Inc.
+# Copyright 2021- Python Language Server Contributors.
+
+import logging
+from pathlib import Path
+
+from pylsp import hookimpl
+from pylsp.lsp import SymbolKind
+
+log = logging.getLogger(__name__)
+
+
+@hookimpl
+def pylsp_workspace_symbol(config, workspace, document, query):
+    if not query or not workspace:
+        return []
+
+    return [_jedi_name_to_symbol(jedi_name) for jedi_name in workspace.search(query)]
+
+
+def _jedi_name_to_symbol(jedi_name):
+    return {
+        "name": jedi_name.name,
+        "kind": _jedi_type_to_symbol_kind(jedi_name.type),
+        "location": {
+            "uri": "file://" + str(jedi_name.module_path),
+            "range": {
+                "start": {"line": jedi_name.line - 1, "character": jedi_name.column},
+                "end": {"line": jedi_name.line - 1, "character": jedi_name.column},
+            },
+        },
+    }
+
+
+def _jedi_type_to_symbol_kind(jedi_type):
+    if jedi_type == "module":
+        return SymbolKind.Module
+    if jedi_type == "class":
+        return SymbolKind.Class
+    if jedi_type == "function":
+        return SymbolKind.Function
+    if jedi_type == "instance":
+        return SymbolKind.Variable
+    if jedi_type == "statement":
+        return SymbolKind.Variable
+    if jedi_type == "param":
+        return SymbolKind.Variable
+    if jedi_type == "path":
+        return SymbolKind.File
+    return SymbolKind.Variable

pylsp/python_lsp.py

@@ -298,6 +298,7 @@ def capabilities(self):
                 "workspaceFolders": {"supported": True, "changeNotifications": True}
             },
             "experimental": merge(self._hook("pylsp_experimental_capabilities")),
+            "workspaceSymbolProvider": True,
         }
         log.info("Server capabilities: %s", server_capabilities)
         return server_capabilities
@@ -433,6 +434,11 @@ def highlight(self, doc_uri, position):
             or None
         )
 
+    def workspace_symbol(self, query):
+        response = self._hook("pylsp_workspace_symbol", query=query)
+        log.debug("Workspace symbol hook returned: %s", response)
+        return flatten(response)
+
     def hover(self, doc_uri, position):
         return self._hook("pylsp_hover", doc_uri, position=position) or {"contents": ""}
 
@@ -767,6 +773,9 @@ def m_text_document__hover(self, textDocument=None, position=None, **_kwargs):
     def m_text_document__document_symbol(self, textDocument=None, **_kwargs):
         return self.document_symbols(textDocument["uri"])
 
+    def m_workspace__symbol(self, textDocument=None, **_kwargs):
+        return self.workspace_symbol(_kwargs["query"])
+
     def m_text_document__formatting(self, textDocument=None, options=None, **_kwargs):
         return self.format_document(textDocument["uri"], options)
 

pylsp/workspace.py

@@ -100,6 +100,46 @@ def root_path(self):
     def root_uri(self):
         return self._root_uri
 
+    def search(self, query):
+        project = self.jedi_project()
+        return project.complete_search(query)
+
+    def jedi_project(self, use_document_path=False):
+        extra_paths = []
+        environment_path = None
+        env_vars = None
+
+        if self._config:
+            jedi_settings = self._config.plugin_settings(
+                "jedi", document_path=self.root_path
+            )
+            jedi.settings.auto_import_modules = jedi_settings.get(
+                "auto_import_modules", DEFAULT_AUTO_IMPORT_MODULES
+            )
+            environment_path = jedi_settings.get("environment")
+            # Jedi itself cannot deal with homedir-relative paths.
+            # On systems, where it is expected, expand the home directory.
+            if environment_path and os.name != "nt":
+                environment_path = os.path.expanduser(environment_path)
+
+            extra_paths = jedi_settings.get("extra_paths") or []
+            env_vars = jedi_settings.get("env_vars")
+
+        # Drop PYTHONPATH from env_vars before creating the environment because that makes
+        # Jedi throw an error.
+        if env_vars is None:
+            env_vars = os.environ.copy()
+        env_vars.pop("PYTHONPATH", None)
+
+        sys_path = extra_paths  # self.sys_path(environment_path, env_vars=env_vars) + extra_paths
+        project_path = self.root_path
+
+        # Extend sys_path with document's path if requested
+        if use_document_path:
+            sys_path += [os.path.normpath(os.path.dirname(self.path))]
+
+        return jedi.Project(path=project_path, sys_path=sys_path)
+
     def is_local(self):
         return (self._root_uri_scheme in ["", "file"]) and os.path.exists(
             self._root_path

pyproject.toml

@@ -71,7 +71,8 @@ jedi_highlight = "pylsp.plugins.highlight"
 jedi_references = "pylsp.plugins.references"
 jedi_rename = "pylsp.plugins.jedi_rename"
 jedi_signature_help = "pylsp.plugins.signature"
-jedi_symbols = "pylsp.plugins.symbols"
+jedi_document_symbols = "pylsp.plugins.document_symbols"
+jedi_workspace_symbol = "pylsp.plugins.workspace_symbol"
 mccabe = "pylsp.plugins.mccabe_lint"
 preload = "pylsp.plugins.preload_imports"
 pycodestyle = "pylsp.plugins.pycodestyle_lint"

test/plugins/test_symbols.py

@@ -8,7 +8,8 @@
 
 from pylsp import uris
 from pylsp.lsp import SymbolKind
-from pylsp.plugins.symbols import pylsp_document_symbols
+from pylsp.plugins.document_symbols import pylsp_document_symbols
+from pylsp.plugins.workspace_symbol import pylsp_workspace_symbol
 from pylsp.workspace import Document
 
 PY2 = sys.version[0] == "2"