Plugin development: Leveraging IDE Python type inference with type hints

A step in our journey to improve Python typing in our code base, helping our developers along the way to read and write type safe code.

A Python Typing 🥁

One year ago, we showed how to add type checking to an existing code base in the article Python Typing with mypy: Progressive Type Checking on a Large Code Base. Since then, we’ve made great strides in adding types to our code base, which both improved the safety of our code and the powerfulness of our tools. For example, our IDEs can now infer types where they couldn’t before. How can we leverage that type inference to help us code faster, safer, stronger?

Looking back at 1 year of typing

Since we introduced mypy type checking 1 year ago, we added many new type hints and a lot of new components to type checking. The power of adding them progressively is that each team can choose to work on typing at their own pace. This is how we distribute ownership at Alan.

The safety upsides of those types are enormous, but as a Kotlin developer and Intellij user, I felt something was missing on the side of developer experience: why can’t I see the inferred types? Why aren’t type hints shown in PyCharm?

Both `one` and `get_subscriptions_for_contracts` function returns are typed, how do I know the type of `subscription`? 🤔

The reason behind this is simply an IDE limitation of PyCharm (see here for the bug report), which doesn’t show the inferred types. But what are inferred types, and why are they useful? What do they eat during winter?

Understanding why type inference is important

Type inference is how type checkers propagates types over functions and variables.

Given a function has the type hint : str

• it returns a value of type str

Given a variable assigned from that method return

• it will be inferred its type is str ➡️ that is type inference
• type inference works transitively for typed functions

This is useful for type checkers, which can track the inferred type of variables and make sure the program is sound over time, which is especially useful for dynamic languages such as Python.

This is also useful for IDEs to propose method calls and variables that are associated with the type:

The type of `my_string` is `str`, so the IDE can propose suggestions accordingly 🤖

Writing a plugin for PyCharm

So PyCharm doesn’t show inferred types, what can we do about it? We can write a plugin!

What we currently have:

The inferred types are not shown 😢

What we want:

The inferred types are shown (e.g. `company: Company`) 😄

Writing a plugin for any JetBrains product looks easy on paper and is well documented, but it is harder than it looks for a multitude of reasons that we’ll get into. Let’s go over the different steps required to make and maintain a plugin.

Setting up the environment

The easiest way to start working on your plugin is to use the IntelliJ Platform Plugin Template, which gives a suggested project structure and defaults. It is also really helpful to get proper versions and links to the documentation.

Once created, your project should look like this: Python Inlay Hints Plugin

Coding the plugin and decompiling code

Depending on your plugin type, the entry point(s) will be different. In our case, we show the type hints using folds, so we need to define a folding entry point.

This is done in resources/META-INF/plugin.xml, under the Extensions / FoldingBuilder tags, where we provide a custom implementation class. In our case this is the PythonVariableTypeHintFolderBuilder.

The code does three things:

(1) iterates on target expressions — which are variable assignments

• ➡️ PsiTreeUtil.findChildrenOfType(r, PyTargetExpression::class.java)

(2) uses the Code Insight API — to get the inferred variable type information

• ➡️ TypeEvalContext.deepCodeInsight(project)

(3) returns a fold — with the variable type information

• ➡️ FoldingDescriptor(…)

The resulting code is relatively short, but took me a lot of time to write. Why is that? It is because PyCharm’s code is not entirely open source, so depending on which part we’re using, we might need to open and debug compiled code.

Reading decompiled code is not fun 😬, but it did help me debug and understand what methods and classes were available, especially for the PSI and Code Insight APIs.

Releasing the plugin to the marketplace

We’re almost there 🎉! We just need to release the plugin to the JetBrains Marketplace and add a description: Python Inlay Hints

Updating the plugin on IDE upgrade

Unfortunately, each version change of the Intellij / PyCharm platform might break the plugin, which happens to us once in a while.

The plugin declares a pluginUntilBuild maximum support version, which triggers a plugin “not compatible with the current version” error when upgrading the IDE.

Most of the time, we only need to bump the supported version and we’re good to go (see here for example). Other times, we need to make changes to API calls, if some of them were removed or changed, but it is less likely to occur.

The dreaded “plugin not compatible” error shown in the user’s IDE 🐛

Leveraging the plugin’s output

Now that we have our plugin released and the engineering team is using it, what are the main takeaways of its usage?

• Readability — We have types for each local and global variables, even in loops and generators, which helps us read and understand the code

The inferred types are shown as type hints in folds (e.g. `parameter: str`) 📂

• Missing or weird types — We can see which objects or functions are not properly typed, by looking at variable with missing or unexpected types

Oh no, this method return type is not properly annotated (so the `constraint_list` variable has no inferred type) 🔥

• Optional — We can see which variables are Optional by looking at the inferred types, meaning we’ll need to handle the None case. This also prompted us to create typing utilities around handling those cases, such as the mandatory helper

The `current_version` is inferred as `Optional`, meaning we have to handle the `None` case 💼

• Hybrid — We can provide both hard coded types and see inferred types at the same time

The `constraint_list` variable has an inferred type from the plugin and the `constraints` variable has hard coded types from the code 💻

• Naming — We can also remove the type from the variable name, since we now see it in the fold, for clearer and more maintainable code. This doubles down as handling your troll colleague as well 😆

🚎

Conclusion

This journey into improving our tooling setup by creating a new plugin has brought us many learnings on our code base and helped us be more productive. At Alan, we have the liberty of using the tools we want, and we are encouraged at improving our setup if it helps our productivity.

What is your setup for handling Python hints and Python type checking in your code base? Let us know on Twitter!