IDE features missing when using PythonVirtualEnvironmentOperator

[cooperellidge]

trafficstars

Apache Airflow version

2.10.0

If "Other Airflow 2 version" selected, which one?

No response

What happened?

I do not get IDE features (LSP stuff) when using PythonVirtualEnvironmentOperator.

from airflow import DAG
from airflow.decorators import task
from airflow.utils.dates import days_ago

# My IDE has Airflow venv activated to get these top-level imports to be recognised.
# After activating a fresh new Airflow venv, I use this to install dependencies
# ```pip install "apache-airflow==${AIRFLOW_VERSION}" --constraint "${CONSTRAINT_URL}" virtualenv```

with DAG(
    "do_something_in_virtual_env",
    start_date=days_ago(1),
) as dag:

    @task.virtualenv(
        task_id="do_something",
        requirements=["sqlalchemy==2.*", "scipy"],
    )
    def do_something() -> None:
        import sqlalchemy  # v1.4.53 is installed in Airflow venv, but I am asking for v2.0.32 in this task's venv
        import scipy # I don't have scipy installed in Airflow venv, but do in this task's venv
        #      ^^^   I get a "Import "scipy" could not be resolved" error

        print(sqlalchemy.__version__) # I have auto-complete but for v1: outputs 2.0.32
        print(scipy.__version__) # I don't have auto-complete: outputs 1.14.1
        #      ^^^   I don't get any IDE features for scipy.

    do_something()

I ran into this use case specifically because a package I require depends on SQLAlchemy >= 2 so I need to run the task in a virtual environment.

What you think should happen instead?

I don't actually know how this could even be solved, as my IDE can only have 1 active virtual environment.

I am mostly wondering what the community's solution to getting IDE features working when using PythonVirtualEnvironmentOperator.

How to reproduce

1. Install Airflow 2.10.0 with constraints file and virtualenv
2. Add the above DAG
3. Run airflow standalone
4. Manually run the DAG

Inspecting the logs shows successfully setup of the virtual environment and the expected package versions being printed. Additionally a pip list in the Airflow virtual environment shows the expected package versions too.

However, my IDE (vscode) cannot provide LSP features in the task definition for the task's virtual environment.

Operating System

Ubuntu

Versions of Apache Airflow Providers

No response

Deployment

Other

Deployment details

Just playing around locally with airflow standalone.

Anything else?

No response

Are you willing to submit PR?

• [ ] Yes I am willing to submit a PR!

Code of Conduct

• [X] I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[phi-friday]

Depending on which IDE you're using, the method will likely be different. I'm using vscode, so I get help from pylance and pyright. In this case, you can get some help from IDE by adding the necessary modules to the typings path.

[cooperellidge]

I'm also using VS Code with Pylance and Mypy.

I have selected the interpreter as my Airflow venv. In this example, that includes Airflow and SQLAlchemy v1. I can create a second virtual env locally that includes scipy and SQLAlchemy v2 with a path /path/to/other/venv.

In my vscode settings.json, I have added:

{
    "python.autoComplete.extraPaths": [
        "/path/to/other/venv/bin/python"
    ],
    "python.analysis.extraPaths": [
       "/path/to/other/venv/bin/python"
    ]
}

I still see the same pylance reportMissingImports error as before on scipy.

And even if I could add scipy typings to my IDE, I don't see how this could resolve the v1 / v2 issue on SQLAlchemy.

[phi-friday]

The extraPaths might be valid when you add the plugin, but it doesn't matter at this point. We just need to make sure the typings are correct. Using typings to use sqlalchemy@v2 is also a workaround, albeit limited. If you're using uv, you could probably try using override (be careful, it will throw an error at runtime).

[cooperellidge]

Can you be more specific?

How do I make the typings correct? What is typings? Is it the typing python package, a VS code setting or configuration, a python virtual environment configuration?

I modified my settings.json paths to point to the custom venv's site-packages rather than the python binary and that helped by allowing my IDE to resolve scipy (in this example). However, it still doesn't resolve the sqlalchmey v1 or v2 issue. And I also have to manually maintain separate virtual environments on my local machine.

[phi-friday]

The typings is the path specified as the default value for the stubPath used by pyright. If you're using pyright or pylance, you can use the module in typings as if you were using it. see more: https://github.com/microsoft/pyright/blob/main/docs/configuration.md#environment-options