Intellisense Support for Airflow Plugins in VSCode - visual-studio-code

Question

I'm trying to find a way to make VSCode Python intellisense work with Airflow Plugins.
Following the code example the import path of plugin operators could be:
from airflow.operators import MyPluginOperator
VSCode cannot resolve this import because it will only be valid at runtime through the airflow plugin system.
Is there any way to configure VSCode to resolve this import?

Answer 1

Airflow loads plugins dynamically by searching the airflow/plugins folder for AirflowPlugin subclasses and add them in airflow namespace in a runtime. Here is the code from airflow/operators/__init__.py:
# Imports operators dynamically while keeping the package API clean,
# abstracting the underlying modules
...
def _integrate_plugins():
"""Integrate plugins to the context"""
from airflow.plugins_manager import operators_modules
for operators_module in operators_modules:
sys.modules[operators_module.__name__] = operators_module
globals()[operators_module._name] = operators_module
VS Code can't handle it. Even "big" Python IDEs like PyCharm has problems with it. It is impossible for VS Code to know that a piece of code in particular folder will transform in airflow.operator later. "python.autoComplete.extraPaths" will not help too. You should only hope that someone will write a VS Code extension for Airflow somewhere :)

Answer 2

Since version 2.0.0 the way that Airflow handles plugins imports has changed:
Importing operators, sensors, hooks added in plugins via airflow.{operators,sensors,hooks}.<plugin_name> is no longer supported, and these extensions should just be imported as regular python modules. For more information, see: Modules Management and Creating a custom Operator
The next important thing to consider is that Airflow adds dags/, plugins/, and config/ directories to PYTHONPATH env. source
Following the specifications from the docs, you could create your MyCustomOperator
in airflow_home/plugins/custom_operator/ directory. Then you could use it like this:
from custom_operator.my_custom_operator import MyCustomOperator
with dag:
sample_custom_task = MyCustomOperator(task_id='sample-task')
So far so good, the dag will run, but according to my experience, the VSCode IntelliSense won't work yet. To make it work, you need to add the path to /plugins as same as Airflow will do when it runs. Depending on your setup there are a few ways to do this. The goal is to add an extra path to the python interpreter VSCode is "using" (make sure to select the interpreter related to your project)
A common solution could be to use the env PYTHONPATH to append our path to the path the interpreter knows. In my case, I'm using a virtual environment, so following the explained in this post, I created a .pth file with the path I wanted to add and locate that file on airflow_home/venv/lib/my_python_version/site-packages/
Following the path on our example, this will create such a file:
cd $(python -c "from distutils.sysconfig import get_python_lib; print(get_python_lib())")
echo airflow_home/plugins > airflow_plugins_example.pth
Once that is done, reload VSCode (could also just change to another interpreter and then come back) and you should see the intelliSense working properly.
Hope that works!