Skip to content

Commit

Permalink
Annotate array_namespace and add py.typed
Browse files Browse the repository at this point in the history
This appears to be the only missing type annotation in the project.
  • Loading branch information
NeilGirdhar committed Feb 20, 2025
1 parent 66106cb commit 1114aa5
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 2 deletions.
4 changes: 2 additions & 2 deletions array_api_compat/common/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

if TYPE_CHECKING:
from typing import Optional, Union, Any
from ._typing import Array, Device
from ._typing import Array, Device, Namespace

import sys
import math
Expand Down Expand Up @@ -439,7 +439,7 @@ def _check_api_version(api_version: str) -> None:
raise ValueError("Only the 2023.12 version of the array API specification is currently supported")


def array_namespace(*xs, api_version=None, use_compat=None):
def array_namespace(*xs, api_version=None, use_compat=None) -> Namespace:
"""
Get the array API compatible namespace for the arrays `xs`.
Expand Down
2 changes: 2 additions & 0 deletions array_api_compat/common/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
"SupportsBufferProtocol",
]

from types import ModuleType
from typing import (
Any,
TypeVar,
Expand All @@ -22,3 +23,4 @@ def __len__(self, /) -> int: ...
Array = Any
Device = Any
DType = Any
Namespace = ModuleType
Empty file added array_api_compat/py.typed
Empty file.
3 changes: 3 additions & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,4 +34,7 @@
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
],
package_data={
"array_api_compat": ["py.typed"],
},
)

0 comments on commit 1114aa5

Please sign in to comment.