support a `Module` type in `ty_extensions` for typing specific modules

[KotlinIsland]

Summary

from ty_extensions import Module


def get_module() ->  Module["collections.abc"]:
    from collections import abc
    
    return abc

this is already possible by using TypeOf (thank you @AlexWaygood):

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import collections.abc
    from ty_extensions import TypeOf

def get_module() -> TypeOf[collections.abc]:
    import collections.abc
    return collections.abc

[AlexWaygood]

Can you give a more full example of a situation in which you might use this? :-) are you thinking that you might use this to annotate functions that lazy-load expensive modules to reduce import times?

[augustfe]

It's maybe not the same use case, but it shows up in the jax library, causing funky decorator behaviour with their @export,

https://github.com/jax-ml/jax/blob/c4481d422bc9c0e4816dd8d895709e9c27be1f8c/jax/_src/nn/initializers.py#L39

export = set_module('jax.nn.initializers')

based on this method

https://github.com/jax-ml/jax/blob/c4481d422bc9c0e4816dd8d895709e9c27be1f8c/jax/_src/util.py#L777-L782

def set_module(module: str) -> Callable[[T], T]:
  def wrapper(func: T) -> T:
    if module is not None:
      func.__module__ = module
    return func
  return wrapper

[AlexWaygood]

@augustfe I'm not sure that's related at all (I think your issue is #1136)

[jorenham]

I like this 👌. And for what it's worth; the pyrights do something similar:

import math
from typing import reveal_type

reveal_type(math)  # Type of "math" is "Module("math")"

https://basedpyright.com/?typeCheckingMode=strict&code=JYWwDg9gTgLgBCAhjAFgKAGZQiOMCeYwAdgOZyiSxxQCmAbrYgDYD6BYtaadjL7hWgAokqAJRw4AYjgAVQXAgY4AIlEoVFAM6qAshAAmAV2bC1yDWJVogA

[AlexWaygood]

I like this 👌. And for what it's worth; the pyrights do something similar:

Right, we already do that too, but I don't believe pyright provides what this feature request is asking for, which is a way of expressing that "module-literal" type in user annotations. (Don't know about basedpyright.)

are you thinking that you might use this to annotate functions that lazy-load expensive modules to reduce import times?

@KotlinIsland, you upvoted this comment, so I assume the answer is "yes". This wouldn't cover all cases where you use functions to lazy-load modules -- e.g. functions that return specific functions or classes from a lazily loaded module -- we would need to provide a way for users to express function-literal types and class-literal types too in order to be able to annotate those functions. But in the meantime, I guess you can always annotate those functions as returning Callable/type[] types (though these are not as precise as function-literal and class-literal types).

[jorenham]

As for a use-case, mpmath comes to mind, which sympy heavily relies on. It uses module-level polymorphism, and they refer to a polymorphic modules as a "context" (docs).
At one point I attempted writing stubs for it, but I wasn't able to get very far because of this. But with the proposed Module type I think this should be possible.

[AlexWaygood]

Ah yeah, you can always annotate module-literal types using Protocols for that kind of thing but that gets painful pretty quickly. That's definitely a reasonable use case.

[KotlinIsland]

Ah yeah, you can always annotate module-literal types using Protocols for that kind of thing but that gets painful pretty quickly. That's definitely a reasonable use case.

yes! i've run into this before, same with lazy loading mentioned above with some painful import cycles

i wonder if there's any value in a special form that sucks the nominality off something, a little like CallableTypeOf, and creates a protocol from it

from ty_extensions import ProtocolOf

class A:
    a: int

class B:
    a: int

type P = ProtocolOf[A]

p: P = B()

[dangotbanned]

Just for another example, in (https://github.com/narwhals-dev/narwhals) we'd really benefit from this since all dependencies are optional.

For some internal things where we can eat the breakage in typing, something like this hack (narwhals-dev/narwhals#2051 (comment)) can kinda work.

But we have a bunch of public functions that return entire packages, which:

1. Would be unrealistic to stub out correctly for even 1 package (let alone 16! 😭)
2. Needs to be expressed as a Module("package") | None, so can't rely on the inference hack we have elsewhere

[AlexWaygood]

Oh, I forgot -- we actually do already have a way of annotating module-literal types; I don't know if this is sufficient?

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import collections.abc
    from ty_extensions import TypeOf

def get_module() -> TypeOf[collections.abc]:
    import collections.abc
    return collections.abc

[KotlinIsland]

well, maybe this isn't required then? probably it would help make things a little easier to understand, but is it worth is, i'm not really sure