3

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?

J.Doe
  • 53
  • 1
  • 6

2 Answers2

2

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 :)

vurmux
  • 9,420
  • 3
  • 25
  • 45
  • Thanks for the explanation. I will try to write a VSCode extension for airflow plugins. – J.Doe Jul 17 '19 at 07:13
  • @J.Doe Have you started the extension? I'd like to have a go at it if not. – Dvir Yitzchaki Aug 27 '19 at 08:27
  • @DvirYitzchaki Not yet! Please let us know if you get it to work. – J.Doe Sep 11 '19 at 06:48
  • 1
    @DvirYitzchaki Hey, I just found this article which seems to work perfectly for hooks/operators: https://www.astronomer.io/guides/airflow-importing-custom-hooks-operators/ – J.Doe Oct 11 '19 at 11:48
  • I keep getting import errors. Could you checkout a question I posted about this? https://stackoverflow.com/questions/64535476/trouble-with-imports-in-vscode-and-pythonpath – Imad Oct 26 '20 at 10:42
1

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!

NicoE
  • 4,373
  • 3
  • 18
  • 33
  • Hi, thank you for the detailed explanation. You can do the same think in Airflow 1.x without problems. I place all my operators in plugin/operators and added a .env file to my vscode root which adds the plugin directory to python path. It seems to be a bit easier oppossed to the .pth file. See this post for reference: https://www.astronomer.io/guides/airflow-importing-custom-hooks-operators – J.Doe Feb 10 '21 at 10:34
  • That's good to know! I haven't tried it on Airflow 1.x. Back when I was trying to figure this out, I did saw the article you mention, but It wasn't clear to me then how to add the plugins directory as a source root to the IDE (it's not detailed on the article). I think the overall key concept is to add the path to /plugins, being the means to do that, `PYTHONPATH` env var. Also, I think is important to notice that opposite to the previous answer, VSCode IntelliSense works just fine with Airflow plugins. – NicoE Feb 10 '21 at 13:46