__array_namespace__ type Hint

[gilfree]

I hope this falls under the scope of this repo -

Is there any a way type-hint __array_namespace__? Is it even possible to do so, and it needs to be standardized, or maybe it's not possible with current python type hinting system?

Motivation:

Suppose I want to write a function that will work in multiple conforming libraries, it will probably start with:

xp = x.__array_namespace__()

or, like in NEP-47:

def foo(x,y):
   xp = get_namespace(x,y)

In both cases, I want to be able to type check xp, and have auto-complete on xp in the various IDEs. (At least some auto-complete engines relay on static typing e.g https://github.com/microsoft/pyright)

Is this possible?

[BvB93]

Unfortunately you can't return Literal modules (which would make this a lot easier), but when implementing the API it would be possible to define and return an auxiliary class that mirrors the content of the namespace.

from __future__ import annotations

import types
import numpy.array_api
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    # Optionally subclass `types.ModuleType` to closer match runtime behavior
    class _ArrayAPINameSpace(types.ModuleType):
        asarray = numpy.array_api.asarray
        arange = numpy.array_api.arange
        empty = numpy.array_api.empty
        # etc
else:
    import numpy.array_api as _ArrayAPINameSpace


class Foo:
    def __array_namespace__(self, /, *, api_version: None | str = None) -> _ArrayAPINameSpace:
        return numpy.array_api

[BvB93]

Small update: unfortunately https://github.com/python/mypy/issues/708 was never completelly fixed, so you'd have to return type[_ArrayAPINameSpace] in the case of mypy, otherwise it will interpret the namespace functions as normal methods and strip their first argument.

class Foo:
    def __array_namespace__(self, /, *, api_version: None | str = None) -> type[_ArrayAPINameSpace]:
        return numpy.array_api

[asmeurer]

This issue is also somewhat related https://github.com/data-apis/array-api/issues/229

[gilfree]

Hi @BvB93, Thanks for the answer.

Are there any thoughts about making such a type/class part of the standard (It may be helpful for statically checking conformance)?

[BvB93]

My personal thoughts: the standard currently specifies that __array_namespace__ should return "an object representing the array API namespace", so concrete API implementations should have enough flexibility to use this little type-checking trick if they'd so desire.

[gilfree]

Thanks @BvB93

[BvB93]

Recently the __getattr__ method to types.ModuleType in typeshed (xref https://github.com/python/typeshed/pull/6302), meaning that all module getattr operations lacking a dedicated set of annotations will return Any. Now, while this is by no means as good as a dedicated type representing the namespaces' content, it does provide a notable improvement over just returning Any (assuming the namespace is actually a module during runtime).

- def __array_namespace__(self, /, *, api_version: None | str = None) -> Any: ...
+ def __array_namespace__(self, /, *, api_version: None | str = None) -> types.ModuleType: ...

[rgommers]

That sounds like a good idea. typeshed isn't a common dependency though for libraries. Is it going to land in typing_extensions? What should we do with this issue, reopen to keep track of the idea?

[BvB93]

That sounds like a good idea. typeshed isn't a common dependency though for libraries. Is it going to land in typing_extensions?

So typeshed is the official repro containing annotations for the standard library, copies of which are vendored by most (all?) type checkers. While you could manually update it, it's general generally more convenient to just wait until mypy and the likes have a new release.

What should we do with this issue, reopen to keep track of the idea?

It might be worthwhile, though I don't expect any major breakthroughs without a mypy plugin of some sort. While comparatively minor, the upstream types.ModuleType change is still a nice bonus.

[rgommers]

Ah okay, got it now - types.ModuleType is already in the stdlib for Python 3.8, so it's fine for us to use it now. And Mypy et al. supporting the new getattr behavior will materialize. So I'd say let's use ModuleType as the annotation for the return type in implementations where the return type is actually a module.

Note that in the standard, the docs now say **out**: _<object>_ and I'm not sure we can do better there (aside from a recommendation), because the returned object doesn't need to be a module, it could be a class instance for example.

[BvB93]

Note that in the standard, the docs now say **out**: _<object>_ and I'm not sure we can do better there (aside from a recommendation), because the returned object doesn't need to be a module, it could be a class instance for example.

I feel that Any would be a small improvement over object here; it is designed as the ultimate placeholder, used for situations situation wherein some specific(-ish) type is involved, but you're not sure how to express it. In contrast, object is generally reserved for situation wherein truly any arbitrary object is considered valid.