Finding References and Exploring the Language-Server Protocol

I spent some time at the end of last year struggling to setup a smooth Python development environment in emacs. If you make a quick search on how to set up a good development environment for Python in Emacs you can't get around the language-server protocol (LSP). The amount of information can also be bewildering if you are someone like me, i.e. a data person who somehow fell into software development by chance and necessity who mainly works on very domain-specific stuff within data science and optimisation. All the information and the choices I had to made quickly became bewildering. What is the protocol doing, what is the client doing and how do I connect it to the server? How do I ensure that the backend can access the correct libraries and project modules?

A quick search for "emacs lsp" leads you to lsp-mode, which is very well-documented, and the Palantir language server (pyls). However, despite the documentation, I couldn't make it run smoothly. It was slow and prone to crashing. Restarting it was also a non-trivial task (at least for me). That led me to the Microsofts Python language server, which worked pretty smoothly whenever it worked properly, but now I had trouble running the language server in the right environment and couldn't figure out why. As opposed to pyls, which is written in Python, the Microsoft language server is written in TypeScript. This meant I couldn't manage it within my virtual environments. At some point I had spent enough time stumbling around blindly and I crossed the threshold where I decided it might make sense to try and understand a bit better how these pieces work together and how seemingly magic features such as renaming across a project, finding references, auto-completion and so on work under the hood.

In the end I landed on a setup that works without any hassle using eglot as the client and pyls as the server. If you somehow ended up here looking for an example configuration you can find it here.

A feature I use very often when exploring new code is find-references, i.e. find all uses of the object at point in the current project. What actually happens when I invoke that command? This is pretty easy to find out in emacs. The call stack can be found by starting the emacs profiler with M-x profiler-start and then checking it after running the command with M-x profiler-report. The result is seen below.

This is a helpful step to figure out where to look inside the eglot source code, which luckily is contained in one 3000 SLOC file. eglot also logs all requests and responses sent from and received by us, which adheres to the JSON-RPC protocol. I use the python-language-server, which relies on jedi to provide auto-completion, finding references, refactorings such as renaming a variable or function, code search and so on.

Let's follow the call. I'll explore the jedi project itself and look for all usages of the Script class, which according to the documentation is the base for completions, goto or whatever you want to do with Jedi. The request sent by eglot to the running server looks like this:

[client-request] (id:18) Tue Jan 26 20:10:38 2021:
(:jsonrpc "2.0" :id 18 :method "textDocument/references" :params
	  (:textDocument
	   (:uri "file:///home/haakon/jedi/jedi/api/__init__.py")
	   :position
	   (:line 50 :character 6)
	   :context
	   (:includeDeclaration t)))

A quick look at the JSON-RPC specification suffices to understand what is going on here. We want to call the "textDocument/references" method with the parameters textDocument, position and context. The document and the position gives us a way to detect the object whose references we are looking for.

The response to the request above looks like this (abbreviated):

[server-reply] (id:18) Tue Jan 26 20:10:38 2021:
(:jsonrpc "2.0" :id 18 :result
	  [(:uri "file:///home/haakon/jedi/conftest.py" :range
		 (:start
		  (:line 109 :character 24)
		  :end
		  (:line 109 :character 30)))
	   (:uri "file:///home/haakon/jedi/conftest.py" :range
		 (:start
		  (:line 115 :character 24)
		  :end
		  (:line 115 :character 30)))
                  <snip>
	   (:uri "file:///home/haakon/jedi/scripts/wx_check.py" :range
		 (:start
		  (:line 47 :character 9)
		  :end
		  (:line 47 :character 15)))
	   (:uri "file:///home/haakon/jedi/sith.py" :range
		 (:start
		  (:line 122 :character 35)
		  :end
		  (:line 122 :character 41)))])

eglot then processes this response and leverages built-in functionality in emacs (xref mainly) in order to show these locations and provide clickable links to their locations, amongst other things. This quick exploration made things a bit more clear: eglot is the client and sends the request to the running backend, which is the python language server. It provides the references and xref takes care of providing a list of clickable links to all the references. However, this all the led me to many more questions. First of all, how does the backend find all the references, or, in other words: how the hell does jedi find references?

It is often valuable to try and attack a problem with a naive and open-minded mind. So let's try and implement a scaled down, quick and dirty version of find-references. The JSONRPC messages showed that the position is a useful input which kicked me off. The most obvious solution is to get the name of the object the cursor is pointing to and search for occurences of that name, using string matching.

Let's start with a small string of source code that we can use to test our implementation of find_references

source = '''
variable = 2

def increment(x):
    x = x + 1

print(variable)
increment(variable)'''

We then parse the string using the same parser as jedi uses, parso, although we could also use the ast module, which is part of Python's standard library.

import parso  # Parser used by jedi
module = parso.parse(source)

After parsing the string we have an abstract syntax tree (AST).

def traverse_parse_tree(ast: parso.python.tree.Module):
     current_leaf = ast.get_first_leaf()
     leaves = [current_leaf]
     while (next_leaf := current_leaf.get_next_leaf()):
	 current_leaf = next_leaf
	 leaves.append(current_leaf)
     return leaves

The AST is a tree, meaning that it can be traversed by any standard tree algorithm, but the AST created by parso's parse method conveniently provides a way of finding all the leaves of the tree. The leaves are basically tokens, which are the smallest meaningful parts of the programming language, such as symbols like +, -, =, (, ) or strings like if and while that (keywords in Python) and object names, such as print or current_leaf. The parser does a lot of the heavy lifting here, but the traversal could be easily implemented with depth-first or breadth-first search.

def find_references(line, column, ast: parso.python.tree.Module):
    name = ast.get_name_of_position((line, column))
    leaves = traverse_parse_tree(ast)
    references = [leaf for leaf in leaves if leaf.type == "name" and leaf.value == name.value]
    for reference in references:
	print("{}:{}\t\t{}".format(reference.line, reference.column, reference.parent.get_code()))

With a list of leaves, which already contains the methods we need, we can easily find all references relying on string matching. We just say that a reference is any other object with the same name.

find_references(8, 11, module)

Which prints this to the stdout

2:0		
variable = 2
7:6		(variable)
8:10		(variable)

Arguably quite ugly, but we see that it works for this simple source. We get line numbers and column numbers that can be used to provide links to the locations and we also print the code corresponding to the parent node, which for the first reference is the whole line, but for the other two the parent node has the leaves (, variable ), which is not very useful context. However, providing the right amount of context is implementation details.

This is code that I wrote relatively recently, slightly simplified, in order to quickly test some code that is used to perform some scheduling.

valid_permutations = []
l = [1, 2, 3, 4, 5, 6]
for k in range(len(l)):
    permutations = list(itertools.permutations(l, k))
    for p in permutations:
	leftovers = [i for i in lst if i not in p]
	leftover_permutations = itertools.permutations(leftovers)
	for l in leftover_permutations:
	    valid_permutation = [list(p), list(l)]
	    valid_permutations.append(initial_solution)

The test failed, but not because the code being tested was wrong, but because I hacked together this input generator a bit too quickly. As you can see l is used several times, but for different things. Python will happily accept this, but in the second round of the outermost loop l is no longer the original list. It took me a while to find this error, but a small coffee break and coming back with fresh eyes made the trick. If we call the shitty find-references on the l name in the innermost loop we will get all occurences of l. This is not correct since they refer to different objects: the innermost l is refers to the temporary name assigned from the iterator, while the other l is the actual list we begin with. This is of course no surprise as the first version just relied on string matching and we gladly ignored any kind of context and any notion of scope. The language server however knows this and a quick call to rename was enough to fix the code. Let's dig a bit deeper into how this is done in practice in jedi.

As we briefly saw above, the Script class is the main point of access to anything jedi can do. The python language server takes care of the dispatch, but in the end the method that is called is the get_references method of the Script class. Going deeper we find jedi.inference.references and find_references, which implements the core logic of finding all usage of the object at point. If we go deeper we end up at the Name class of parso which implements a get_definition function. This immediately leads to the next idea: all names that have the same definition as the current object at point are references.

To find references we only need climb up the tree until find the definition of the object. Everything else is searching through the project (modules) to find names whose definition is the same. This also takes care of scoping, e.g. in the example above the definition of the innermost l is actually in the for loop expression, which is the first expression we find when climbing up the AST looking for expressions that define new names.

We now have a better idea of how everything works so let's make our shitty find_references slightly less shitty. Let's use a similar code as the one that broke the first version of find_references.

import parso
import copy
import jedi

source = '''
import os

l = [1, 2, 3, 4]
for k in range(3):
    combos = itertools.combinations(l, k)
    for l in combos:
	print(l)'''

module = parso.parse(source)

The new find_references uses this idea. Instead of string matching we check whether the definition is the same.

def find_references(line, column, ast: parso.python.tree.Module):
    name = ast.get_name_of_position((line, column))
    definition = get_definition(name, ast)
    leaves = traverse_parse_tree(ast)
    def _same_definition(a, b):
	return True if (a.line == b.line and a.value == b.value) else False
    names = [leaf for leaf in leaves if leaf.type == "name"]
    defs = [get_definition(name, ast) for name in names]
    matches = [name for name, d in zip(names, defs) if d is not None and _same_definition(d, definition)]
    for match in matches:
	print("{}:{}\t\t{}".format(match.line, match.column, match.parent.get_code()))

The new version relies on the function get_definition, which checks whether the node itself is a definition, e.g. l = [1, 2, 3, 4], and if not gets all definitions, filters them for the definitions that has the same value and uses a hack to set the actual definition as the last definition that has a smaller line number in the source code.

DEF_TYPES = ["for_stmt", "expr_stmt"]
def get_all_definitions(ast):
    """Parse the tree and grab all definition statements"""
    definitions = []
    if hasattr(ast, "children"):
	for child in ast.children:
	    if child.type in DEF_TYPES:
		definitions.append(child)
	    definitions.extend(get_all_definitions(child))
	return definitions
    else:
	return []

def clean_definitions(definitions):
    cleaned_defs = []
    for d in definitions:
	if d.type == "for_stmt":
	    cleaned_defs.append(d.children[1])
	else:
	    cleaned_defs.append(d.children[0])
    return cleaned_defs

def get_definition(node, ast):
    definition = node.get_definition()
    if definition is not None:
	return node
    else:
	cleaned_defs = clean_definitions(
	    get_all_definitions(ast))
	relevant_defs = [d for d in cleaned_defs if d.value == node.value]
	definition = [d for d in relevant_defs if d.line < node.line]
	return definition[-1] if definition else None

We fixed the shitty version which relied purely on string matching with a version that traverses the abstract syntax tree and returns all uses of a name whose definition is the same and we now only get the actual references.

We obviously simplified a way a lot of things here and hacked our way to a working solution. To make a library like jedi usable it has to be performant as no one is going to wait ages for auto-completion or finding references. It has to work with projects. I gladly skipped that part that although it brings its own challenges with traversing the project to find relevant modules and relevant definitions. Lastly, it has to deal with the whole grammar, which is only getting larger as new features are added (such as the := operator in Python 3.8). I only dealt with assignments via for statements and expr statements, but there are many more.

As we have seen, finding references can be reduced to different traversing a project and the ASTs for each module. Renaming leverages finding referances and just manipulates the AST. jedi is in the end a well-structured library that is able deal with a lot of the complexity that arises from Python's grammar and dynamic nature fast enough. However, since there is no magic that means that it can easily be broken using methods such as setattr or just using multiple labels.

class A:
    pass

a = A()

f = lambda x, y: x + y

setattr(a, "add", f)
a.add(2, 3)

Doing find_references on f won't give you a.square and you won't be able to auto-complete a.a.. or provide automatic documentation as this method only exists at runtime.

In the end I have a much better idea of how everything works together although I spent more time on jedi than I expected. Maybe studying pyls and eglot will be another part. However, it is very nice to see how four different pieces of software (xref, pyls, jedi, eglot) can work smoothly together through well-designed interfaces and protocols.