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

Start with the Proposal and the linked NumPy, lazy-import, and version-dependent import examples, then review the existing typing issue discussion. No implementation file or test is named; done means the supported way to express a module-valued return and its auto-completion and type-checking behavior is specified.