stubgen: import `Union` if generated stub code includes

[eggplants]

Feature

I would like to insert from typing import Union when generated stub code by stubgen includes Union.

Pitch

ref: #12920

Now:

# Python 3.10.4, mypy 0.960
$ cat test.py
import random


def a() -> int | str:
    if random.random() > 0.5:
        return 1
    else:
        return "hey!"

def b() -> str | None:
    if random.random() > 0.5:
        return "hey!"

$ stubgen test.py
Processed 1 modules
Generated out/test.pyi

$ cat out/test.pyi
def a() -> Union[int, str]: ...
def b() -> Union[str, None]: ...

Expect:

from typing import Union

def a() -> Union[int, str]: ...
def b() -> Union[str, None]: ...

[JelleZijlstra]

Thanks! I expect this is an easy fix (just need to ad some function call to make stubgen emit the import), so if you're willing to dive into the code, feel free to make a PR.

[eggplants]

OK, I'll try it...

[bo-dani]

Is there any update on this fix?

[sguidotti]

If no one else is working on this I can pick it up

[eggplants]

@sguidotti

@hamdanal has already created: #14601