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

Bắt đầu với Proposal và các ví dụ import NumPy, lazy-import và import phụ thuộc phiên bản được liên kết, sau đó xem lại phần thảo luận về typing issue hiện có. Không có tệp triển khai hay bài kiểm thử nào được nêu tên; được xem là hoàn tất khi cách được hỗ trợ để biểu đạt giá trị trả về là một module cùng với hành vi tự động hoàn thành và kiểm tra kiểu của nó được đặc tả.