Annotate function which return a specific module (e.g. `-> Literal[np]`)

Problem

Currently it's not possible to annotate a function which return a specific module, like:

np = load_numpy()  # def load_numpy() -> ??:

x = np.array(123)  # << I want to have auto-complete & type checking here

• -> types.ModuleType does not work as too generic (no auto-complete, nor static type checking,...)
• Protocol is not applicable in practice: maintaining a numpy protocol which has 500+ symbols is just not realistic

Use case

Here are some concrete use-case where this feature is needed (also collected from this thread comments):

1. Numpy Enancement Propopsal 37 propose a recipe to have code working with various numpy implementation ( numpy, jax.numpy, tensorflow.numpy):

   def duckarray_add_random(array):
       module = np.get_array_module(array)           # def get_array_module() -> Literal[np]
       noise = module.random.randn(*array.shape)     # << I want to have auto-complete & type checking here
       return array + noise
   

   I developed my version at https://github.com/google/etils/tree/main/etils/enp#code-that-works-with-nparray-jnparray-tftensor

2. Lazy imports is a common pattern to only import a module if needed. Like: https://github.com/tensorflow/datasets/blob/76f8591def26afaca16340b06d057553582f6163/tensorflow_datasets/core/lazy_imports_lib.py#L40-L197

   beam = lazy_import.apache_beam
   
   beam.Pipeline()  # << No auto-completion
   

3. From another user comment: Similar issue to to 1. encountered at: https://github.com/data-apis/array-api/issues/267

4. From another user comment:

   I encountered a similar issue before. Although not with with the return type but rather version dependent imports, e.g. assign either ast (Py >= 3.8) or typed_ast.ast3 to a common variable.

Proposal

I would like to annotate my function as:

def load_numpy() -> Literal[np]:

Or:

def load_numpy() -> np:

For the lazy-loading case, typing.TYPE_CHECKING pattern could be used:

if typing.TYPE_CHECKING:
  import numpy as np


def load_numpy() -> Literal[np]:
  import numpy as np
  return np

Proposal と、リンクされている NumPy、lazy-import、バージョン依存の import の例から始め、その後、既存の typing issue の議論を確認してください。実装ファイルやテストは指定されていません。モジュールを値とする戻り値を表現するサポート対象の方法と、その自動補完および型チェックの動作が仕様化されていれば完了です。