1# Copyright 2017-2020 Palantir Technologies, Inc.
2# Copyright 2021- Python Language Server Contributors.
3
4from functools import partial
5import logging
6import os
7import socketserver
8import threading
9
10from pylsp_jsonrpc.dispatchers import MethodDispatcher
11from pylsp_jsonrpc.endpoint import Endpoint
12from pylsp_jsonrpc.streams import JsonRpcStreamReader, JsonRpcStreamWriter
13
14from . import lsp, _utils, uris
15from .config import config
16from .workspace import Workspace
17from ._version import __version__
18
19log = logging.getLogger(__name__)
20
21
22LINT_DEBOUNCE_S = 0.5  # 500 ms
23PARENT_PROCESS_WATCH_INTERVAL = 10  # 10 s
24MAX_WORKERS = 64
25PYTHON_FILE_EXTENSIONS = ('.py', '.pyi')
26CONFIG_FILEs = ('pycodestyle.cfg', 'setup.cfg', 'tox.ini', '.flake8')
27
28
29class _StreamHandlerWrapper(socketserver.StreamRequestHandler):
30    """A wrapper class that is used to construct a custom handler class."""
31
32    delegate = None
33
34    def setup(self):
35        super().setup()
36        self.delegate = self.DELEGATE_CLASS(self.rfile, self.wfile)
37
38    def handle(self):
39        try:
40            self.delegate.start()
41        except OSError as e:
42            if os.name == 'nt':
43                # Catch and pass on ConnectionResetError when parent process
44                # dies
45                # pylint: disable=no-member, undefined-variable
46                if isinstance(e, WindowsError) and e.winerror == 10054:
47                    pass
48
49        self.SHUTDOWN_CALL()
50
51
52def start_tcp_lang_server(bind_addr, port, check_parent_process, handler_class):
53    if not issubclass(handler_class, PythonLSPServer):
54        raise ValueError('Handler class must be an instance of PythonLSPServer')
55
56    def shutdown_server(check_parent_process, *args):
57        # pylint: disable=unused-argument
58        if check_parent_process:
59            log.debug('Shutting down server')
60            # Shutdown call must be done on a thread, to prevent deadlocks
61            stop_thread = threading.Thread(target=server.shutdown)
62            stop_thread.start()
63
64    # Construct a custom wrapper class around the user's handler_class
65    wrapper_class = type(
66        handler_class.__name__ + 'Handler',
67        (_StreamHandlerWrapper,),
68        {'DELEGATE_CLASS': partial(handler_class,
69                                   check_parent_process=check_parent_process),
70         'SHUTDOWN_CALL': partial(shutdown_server, check_parent_process)}
71    )
72
73    server = socketserver.TCPServer((bind_addr, port), wrapper_class, bind_and_activate=False)
74    server.allow_reuse_address = True
75
76    try:
77        server.server_bind()
78        server.server_activate()
79        log.info('Serving %s on (%s, %s)', handler_class.__name__, bind_addr, port)
80        server.serve_forever()
81    finally:
82        log.info('Shutting down')
83        server.server_close()
84
85
86def start_io_lang_server(rfile, wfile, check_parent_process, handler_class):
87    if not issubclass(handler_class, PythonLSPServer):
88        raise ValueError('Handler class must be an instance of PythonLSPServer')
89    log.info('Starting %s IO language server', handler_class.__name__)
90    server = handler_class(rfile, wfile, check_parent_process)
91    server.start()
92
93
94class PythonLSPServer(MethodDispatcher):
95    """ Implementation of the Microsoft VSCode Language Server Protocol
96    https://github.com/Microsoft/language-server-protocol/blob/master/versions/protocol-1-x.md
97    """
98
99    # pylint: disable=too-many-public-methods,redefined-builtin
100
101    def __init__(self, rx, tx, check_parent_process=False):
102        self.workspace = None
103        self.config = None
104        self.root_uri = None
105        self.watching_thread = None
106        self.workspaces = {}
107        self.uri_workspace_mapper = {}
108
109        self._jsonrpc_stream_reader = JsonRpcStreamReader(rx)
110        self._jsonrpc_stream_writer = JsonRpcStreamWriter(tx)
111        self._check_parent_process = check_parent_process
112        self._endpoint = Endpoint(self, self._jsonrpc_stream_writer.write, max_workers=MAX_WORKERS)
113        self._dispatchers = []
114        self._shutdown = False
115
116    def start(self):
117        """Entry point for the server."""
118        self._jsonrpc_stream_reader.listen(self._endpoint.consume)
119
120    def __getitem__(self, item):
121        """Override getitem to fallback through multiple dispatchers."""
122        if self._shutdown and item != 'exit':
123            # exit is the only allowed method during shutdown
124            log.debug("Ignoring non-exit method during shutdown: %s", item)
125            raise KeyError
126
127        try:
128            return super().__getitem__(item)
129        except KeyError:
130            # Fallback through extra dispatchers
131            for dispatcher in self._dispatchers:
132                try:
133                    return dispatcher[item]
134                except KeyError:
135                    continue
136
137        raise KeyError()
138
139    def m_shutdown(self, **_kwargs):
140        self._shutdown = True
141
142    def m_exit(self, **_kwargs):
143        self._endpoint.shutdown()
144        self._jsonrpc_stream_reader.close()
145        self._jsonrpc_stream_writer.close()
146
147    def _match_uri_to_workspace(self, uri):
148        workspace_uri = _utils.match_uri_to_workspace(uri, self.workspaces)
149        return self.workspaces.get(workspace_uri, self.workspace)
150
151    def _hook(self, hook_name, doc_uri=None, **kwargs):
152        """Calls hook_name and returns a list of results from all registered handlers"""
153        workspace = self._match_uri_to_workspace(doc_uri)
154        doc = workspace.get_document(doc_uri) if doc_uri else None
155        hook_handlers = self.config.plugin_manager.subset_hook_caller(hook_name, self.config.disabled_plugins)
156        return hook_handlers(config=self.config, workspace=workspace, document=doc, **kwargs)
157
158    def capabilities(self):
159        server_capabilities = {
160            'codeActionProvider': True,
161            'codeLensProvider': {
162                'resolveProvider': False,  # We may need to make this configurable
163            },
164            'completionProvider': {
165                'resolveProvider': True,  # We could know everything ahead of time, but this takes time to transfer
166                'triggerCharacters': ['.'],
167            },
168            'documentFormattingProvider': True,
169            'documentHighlightProvider': True,
170            'documentRangeFormattingProvider': True,
171            'documentSymbolProvider': True,
172            'definitionProvider': True,
173            'executeCommandProvider': {
174                'commands': flatten(self._hook('pylsp_commands'))
175            },
176            'hoverProvider': True,
177            'referencesProvider': True,
178            'renameProvider': True,
179            'foldingRangeProvider': True,
180            'signatureHelpProvider': {
181                'triggerCharacters': ['(', ',', '=']
182            },
183            'textDocumentSync': {
184                'change': lsp.TextDocumentSyncKind.INCREMENTAL,
185                'save': {
186                    'includeText': True,
187                },
188                'openClose': True,
189            },
190            'workspace': {
191                'workspaceFolders': {
192                    'supported': True,
193                    'changeNotifications': True
194                }
195            },
196            'experimental': merge(
197                self._hook('pylsp_experimental_capabilities'))
198        }
199        log.info('Server capabilities: %s', server_capabilities)
200        return server_capabilities
201
202    def m_initialize(self, processId=None, rootUri=None, rootPath=None,
203                     initializationOptions=None, workspaceFolders=None, **_kwargs):
204        log.debug('Language server initialized with %s %s %s %s', processId, rootUri, rootPath, initializationOptions)
205        if rootUri is None:
206            rootUri = uris.from_fs_path(rootPath) if rootPath is not None else ''
207
208        self.workspaces.pop(self.root_uri, None)
209        self.root_uri = rootUri
210        self.config = config.Config(rootUri, initializationOptions or {},
211                                    processId, _kwargs.get('capabilities', {}))
212        self.workspace = Workspace(rootUri, self._endpoint, self.config)
213        self.workspaces[rootUri] = self.workspace
214        if workspaceFolders:
215            for folder in workspaceFolders:
216                uri = folder['uri']
217                if uri == rootUri:
218                    # Already created
219                    continue
220                workspace_config = config.Config(
221                    uri, self.config._init_opts,
222                    self.config._process_id, self.config._capabilities)
223                workspace_config.update(self.config._settings)
224                self.workspaces[uri] = Workspace(
225                    uri, self._endpoint, workspace_config)
226
227        self._dispatchers = self._hook('pylsp_dispatchers')
228        self._hook('pylsp_initialize')
229
230        if self._check_parent_process and processId is not None and self.watching_thread is None:
231            def watch_parent_process(pid):
232                # exit when the given pid is not alive
233                if not _utils.is_process_alive(pid):
234                    log.info("parent process %s is not alive, exiting!", pid)
235                    self.m_exit()
236                else:
237                    threading.Timer(PARENT_PROCESS_WATCH_INTERVAL, watch_parent_process, args=[pid]).start()
238
239            self.watching_thread = threading.Thread(target=watch_parent_process, args=(processId,))
240            self.watching_thread.daemon = True
241            self.watching_thread.start()
242        # Get our capabilities
243        return {
244            'capabilities': self.capabilities(),
245            'serverInfo': {
246                'name': 'pylsp',
247                'version': __version__,
248            },
249        }
250
251    def m_initialized(self, **_kwargs):
252        self._hook('pylsp_initialized')
253
254    def code_actions(self, doc_uri, range, context):
255        return flatten(self._hook('pylsp_code_actions', doc_uri, range=range, context=context))
256
257    def code_lens(self, doc_uri):
258        return flatten(self._hook('pylsp_code_lens', doc_uri))
259
260    def completions(self, doc_uri, position):
261        completions = self._hook('pylsp_completions', doc_uri, position=position)
262        return {
263            'isIncomplete': False,
264            'items': flatten(completions)
265        }
266
267    def completion_item_resolve(self, completion_item):
268        doc_uri = completion_item.get('data', {}).get('doc_uri', None)
269        return self._hook('pylsp_completion_item_resolve', doc_uri, completion_item=completion_item)
270
271    def definitions(self, doc_uri, position):
272        return flatten(self._hook('pylsp_definitions', doc_uri, position=position))
273
274    def document_symbols(self, doc_uri):
275        return flatten(self._hook('pylsp_document_symbols', doc_uri))
276
277    def execute_command(self, command, arguments):
278        return self._hook('pylsp_execute_command', command=command, arguments=arguments)
279
280    def format_document(self, doc_uri):
281        return self._hook('pylsp_format_document', doc_uri)
282
283    def format_range(self, doc_uri, range):
284        return self._hook('pylsp_format_range', doc_uri, range=range)
285
286    def highlight(self, doc_uri, position):
287        return flatten(self._hook('pylsp_document_highlight', doc_uri, position=position)) or None
288
289    def hover(self, doc_uri, position):
290        return self._hook('pylsp_hover', doc_uri, position=position) or {'contents': ''}
291
292    @_utils.debounce(LINT_DEBOUNCE_S, keyed_by='doc_uri')
293    def lint(self, doc_uri, is_saved):
294        # Since we're debounced, the document may no longer be open
295        workspace = self._match_uri_to_workspace(doc_uri)
296        if doc_uri in workspace.documents:
297            workspace.publish_diagnostics(
298                doc_uri,
299                flatten(self._hook('pylsp_lint', doc_uri, is_saved=is_saved))
300            )
301
302    def references(self, doc_uri, position, exclude_declaration):
303        return flatten(self._hook(
304            'pylsp_references', doc_uri, position=position,
305            exclude_declaration=exclude_declaration
306        ))
307
308    def rename(self, doc_uri, position, new_name):
309        return self._hook('pylsp_rename', doc_uri, position=position, new_name=new_name)
310
311    def signature_help(self, doc_uri, position):
312        return self._hook('pylsp_signature_help', doc_uri, position=position)
313
314    def folding(self, doc_uri):
315        return flatten(self._hook('pylsp_folding_range', doc_uri))
316
317    def m_completion_item__resolve(self, **completionItem):
318        return self.completion_item_resolve(completionItem)
319
320    def m_text_document__did_close(self, textDocument=None, **_kwargs):
321        workspace = self._match_uri_to_workspace(textDocument['uri'])
322        workspace.rm_document(textDocument['uri'])
323
324    def m_text_document__did_open(self, textDocument=None, **_kwargs):
325        workspace = self._match_uri_to_workspace(textDocument['uri'])
326        workspace.put_document(textDocument['uri'], textDocument['text'], version=textDocument.get('version'))
327        self._hook('pylsp_document_did_open', textDocument['uri'])
328        self.lint(textDocument['uri'], is_saved=True)
329
330    def m_text_document__did_change(self, contentChanges=None, textDocument=None, **_kwargs):
331        workspace = self._match_uri_to_workspace(textDocument['uri'])
332        for change in contentChanges:
333            workspace.update_document(
334                textDocument['uri'],
335                change,
336                version=textDocument.get('version')
337            )
338        self.lint(textDocument['uri'], is_saved=False)
339
340    def m_text_document__did_save(self, textDocument=None, **_kwargs):
341        self.lint(textDocument['uri'], is_saved=True)
342
343    def m_text_document__code_action(self, textDocument=None, range=None, context=None, **_kwargs):
344        return self.code_actions(textDocument['uri'], range, context)
345
346    def m_text_document__code_lens(self, textDocument=None, **_kwargs):
347        return self.code_lens(textDocument['uri'])
348
349    def m_text_document__completion(self, textDocument=None, position=None, **_kwargs):
350        return self.completions(textDocument['uri'], position)
351
352    def m_text_document__definition(self, textDocument=None, position=None, **_kwargs):
353        return self.definitions(textDocument['uri'], position)
354
355    def m_text_document__document_highlight(self, textDocument=None, position=None, **_kwargs):
356        return self.highlight(textDocument['uri'], position)
357
358    def m_text_document__hover(self, textDocument=None, position=None, **_kwargs):
359        return self.hover(textDocument['uri'], position)
360
361    def m_text_document__document_symbol(self, textDocument=None, **_kwargs):
362        return self.document_symbols(textDocument['uri'])
363
364    def m_text_document__formatting(self, textDocument=None, _options=None, **_kwargs):
365        # For now we're ignoring formatting options.
366        return self.format_document(textDocument['uri'])
367
368    def m_text_document__rename(self, textDocument=None, position=None, newName=None, **_kwargs):
369        return self.rename(textDocument['uri'], position, newName)
370
371    def m_text_document__folding_range(self, textDocument=None, **_kwargs):
372        return self.folding(textDocument['uri'])
373
374    def m_text_document__range_formatting(self, textDocument=None, range=None, _options=None, **_kwargs):
375        # Again, we'll ignore formatting options for now.
376        return self.format_range(textDocument['uri'], range)
377
378    def m_text_document__references(self, textDocument=None, position=None, context=None, **_kwargs):
379        exclude_declaration = not context['includeDeclaration']
380        return self.references(textDocument['uri'], position, exclude_declaration)
381
382    def m_text_document__signature_help(self, textDocument=None, position=None, **_kwargs):
383        return self.signature_help(textDocument['uri'], position)
384
385    def m_workspace__did_change_configuration(self, settings=None):
386        if self.config is not None:
387            self.config.update((settings or {}).get('pylsp', {}))
388        for workspace in self.workspaces.values():
389            workspace.update_config(settings)
390            for doc_uri in workspace.documents:
391                self.lint(doc_uri, is_saved=False)
392
393    def m_workspace__did_change_workspace_folders(self, event=None, **_kwargs):  # pylint: disable=too-many-locals
394        if event is None:
395            return
396        added = event.get('added', [])
397        removed = event.get('removed', [])
398
399        for removed_info in removed:
400            if 'uri' in removed_info:
401                removed_uri = removed_info['uri']
402                self.workspaces.pop(removed_uri, None)
403
404        for added_info in added:
405            if 'uri' in added_info:
406                added_uri = added_info['uri']
407                workspace_config = config.Config(
408                    added_uri, self.config._init_opts,
409                    self.config._process_id, self.config._capabilities)
410                workspace_config.update(self.config._settings)
411                self.workspaces[added_uri] = Workspace(
412                    added_uri, self._endpoint, workspace_config)
413
414        root_workspace_removed = any(removed_info['uri'] == self.root_uri for removed_info in removed)
415        workspace_added = len(added) > 0 and 'uri' in added[0]
416        if root_workspace_removed and workspace_added:
417            added_uri = added[0]['uri']
418            self.root_uri = added_uri
419            new_root_workspace = self.workspaces[added_uri]
420            self.config = new_root_workspace._config
421            self.workspace = new_root_workspace
422        elif root_workspace_removed:
423            # NOTE: Removing the root workspace can only happen when the server
424            # is closed, thus the else condition of this if can never happen.
425            if self.workspaces:
426                log.debug('Root workspace deleted!')
427                available_workspaces = sorted(self.workspaces)
428                first_workspace = available_workspaces[0]
429                new_root_workspace = self.workspaces[first_workspace]
430                self.root_uri = first_workspace
431                self.config = new_root_workspace._config
432                self.workspace = new_root_workspace
433
434        # Migrate documents that are on the root workspace and have a better
435        # match now
436        doc_uris = list(self.workspace._docs.keys())
437        for uri in doc_uris:
438            doc = self.workspace._docs.pop(uri)
439            new_workspace = self._match_uri_to_workspace(uri)
440            new_workspace._docs[uri] = doc
441
442    def m_workspace__did_change_watched_files(self, changes=None, **_kwargs):
443        changed_py_files = set()
444        config_changed = False
445        for d in (changes or []):
446            if d['uri'].endswith(PYTHON_FILE_EXTENSIONS):
447                changed_py_files.add(d['uri'])
448            elif d['uri'].endswith(CONFIG_FILEs):
449                config_changed = True
450
451        if config_changed:
452            self.config.settings.cache_clear()
453        elif not changed_py_files:
454            # Only externally changed python files and lint configs may result in changed diagnostics.
455            return
456
457        for workspace in self.workspaces.values():
458            for doc_uri in workspace.documents:
459                # Changes in doc_uri are already handled by m_text_document__did_save
460                if doc_uri not in changed_py_files:
461                    self.lint(doc_uri, is_saved=False)
462
463    def m_workspace__execute_command(self, command=None, arguments=None):
464        return self.execute_command(command, arguments)
465
466
467def flatten(list_of_lists):
468    return [item for lst in list_of_lists for item in lst]
469
470
471def merge(list_of_dicts):
472    return {k: v for dictionary in list_of_dicts for k, v in dictionary.items()}
473