From 1114aa59e34eb2255b5b413a815218211a9d9b07 Mon Sep 17 00:00:00 2001
From: Neil Girdhar <mistersheik@gmail.com>
Date: Thu, 20 Feb 2025 15:59:35 -0500
Subject: [PATCH] Annotate array_namespace and add py.typed
This appears to be the only missing type annotation in the project.
---
array_api_compat/common/_helpers.py | 4 ++--
array_api_compat/common/_typing.py | 2 ++
array_api_compat/py.typed | 0
setup.py | 3 +++
4 files changed, 7 insertions(+), 2 deletions(-)
create mode 100644 array_api_compat/py.typed
diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py
index 434e7d87..80239c06 100644
--- a/array_api_compat/common/_helpers.py
+++ b/array_api_compat/common/_helpers.py
@@ -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
@@ -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`.
diff --git a/array_api_compat/common/_typing.py b/array_api_compat/common/_typing.py
index 1f916cd8..d8acdef7 100644
--- a/array_api_compat/common/_typing.py
+++ b/array_api_compat/common/_typing.py
@@ -5,6 +5,7 @@
"SupportsBufferProtocol",
]
+from types import ModuleType
from typing import (
Any,
TypeVar,
@@ -22,3 +23,4 @@ def __len__(self, /) -> int: ...
Array = Any
Device = Any
DType = Any
+Namespace = ModuleType
diff --git a/array_api_compat/py.typed b/array_api_compat/py.typed
new file mode 100644
index 00000000..e69de29b
diff --git a/setup.py b/setup.py
index d0a28404..2368ccc4 100644
--- a/setup.py
+++ b/setup.py
@@ -34,4 +34,7 @@
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
],
+ package_data={
+ "array_api_compat": ["py.typed"],
+ },
)