numpy protocols improvements

Author: OutSquareCapital

_duckdb-stubs/_typing.pyi

@@ -10,28 +10,46 @@ if TYPE_CHECKING:
     from ._expression import Expression
     from ._sqltypes import DuckDBPyType
 
+# Numpy protocols
+
 _T_co = TypeVar("_T_co", covariant=True)
 _S_co = TypeVar("_S_co", bound=tuple[Any, ...], covariant=True)
 _D_co = TypeVar("_D_co", covariant=True)
 
-class NPTypeLike(Protocol, Generic[_T_co]): ...
+class NPProtocol(Protocol):
+    """Base Protocol for numpy objects."""
+    @property
+    def dtype(self) -> Any: ...
+    @property
+    def ndim(self) -> int: ...
+    def __array__(self, *args: Any, **kwargs: Any) -> Any: ...
+    def __array_wrap__(self, *args: Any, **kwargs: Any) -> Any: ...
+    @property
+    def __array_interface__(self) -> dict[str, Any]: ...
+    @property
+    def __array_priority__(self) -> float: ...
 
-class NPArrayLike(Protocol, Generic[_S_co, _D_co]):
+class NPScalarTypeLike(NPProtocol, Protocol):
+    @property
+    def itemsize(self) -> int: ...
+
+class NPArrayLike(NPProtocol, Generic[_S_co, _D_co], Protocol):
+    """`numpy.ndarray` Protocol.
+
+    This is needed to accept numpy arrays as literals in expressions, without emitting type checker errors about unknown symbol if the user doesn't have numpy installed.
+
+    Note:
+        Using `np.typing.NDArray` is still the best option for return types.
+    """
     def __len__(self) -> int: ...
     def __contains__(self, value: object, /) -> bool: ...
     def __iter__(self) -> Iterator[_D_co]: ...
-    def __array__(self, *args: Any, **kwargs: Any) -> Any: ...
     def __array_finalize__(self, *args: Any, **kwargs: Any) -> None: ...
-    def __array_wrap__(self, *args: Any, **kwargs: Any) -> Any: ...
     def __getitem__(self, *args: Any, **kwargs: Any) -> Any: ...
     def __setitem__(self, *args: Any, **kwargs: Any) -> None: ...
     @property
     def shape(self) -> _S_co: ...
     @property
-    def dtype(self) -> Any: ...
-    @property
-    def ndim(self) -> int: ...
-    @property
     def size(self) -> int: ...
 
 NumericLiteral: TypeAlias = int | float | Decimal