ENH: improve typing of array creation routines

Add overloads for when a dtype is passed (versus just dtype-like) and
handle the default float64 dtype for ones/zeros.

Author: person142

numpy-stubs/__init__.pyi

@@ -521,9 +521,27 @@ def array(
     subok: bool = ...,
     ndmin: int = ...,
 ) -> ndarray: ...
+@overload
+def zeros(shape: _ShapeLike) -> ndarray[float64]: ...
+@overload
+def zeros(shape: _ShapeLike, *, order: Optional[str] = ...) -> ndarray[float64]: ...
+@overload
+def zeros(
+    shape: _ShapeLike, dtype: Type[_ArbitraryDtype] = ..., order: Optional[str] = ...
+) -> ndarray[_ArbitraryDtype]: ...
+@overload
 def zeros(
     shape: _ShapeLike, dtype: _DtypeLike = ..., order: Optional[str] = ...
 ) -> ndarray: ...
+@overload
+def ones(shape: _ShapeLike) -> ndarray[float64]: ...
+@overload
+def ones(shape: _ShapeLike, *, order: Optional[str] = ...) -> ndarray[float64]: ...
+@overload
+def ones(
+    shape: _ShapeLike, dtype: Type[_ArbitraryDtype] = ..., order: Optional[str] = ...
+) -> ndarray[_ArbitraryDtype]: ...
+@overload
 def ones(
     shape: _ShapeLike, dtype: _DtypeLike = ..., order: Optional[str] = ...
 ) -> ndarray: ...

tests/fail/simple.py

@@ -4,7 +4,7 @@
 
 # Array creation routines checks
 np.zeros("test")  # E: incompatible type
-np.zeros()  # E: Too few arguments
+np.zeros()  # E: All overload variants of "zeros" require at least one argument
 
 np.ones("test")  # E: incompatible type
-np.ones()  # E: Too few arguments
+np.ones()  # E: All overload variants of "ones" require at least one argument

tests/reveal/ndarray_conversion.py

@@ -2,9 +2,6 @@
 
 nd = np.array([[1, 2], [3, 4]], dtype=np.int64)
 
-# dtype of the array
-reveal_type(nd)  # E: numpy.ndarray[numpy.int64*]
-
 # item
 reveal_type(nd.item())  # E: Any
 reveal_type(nd.item(1))  # E: Any

tests/reveal/ndarray_creation.py

@@ -0,0 +1,9 @@
+import numpy as np
+
+reveal_type(np.array([[1, 2], [3, 4]], dtype=np.int64))  # E: numpy.ndarray[numpy.int64*]
+reveal_type(np.zeros((3, 3)))  # E: numpy.ndarray[numpy.float64]
+reveal_type(np.zeros((3, 3), dtype=np.int64))  # E: numpy.ndarray[numpy.int64*]
+reveal_type(np.zeros((3, 3), order='F'))  # E: numpy.ndarray[numpy.float64]
+reveal_type(np.ones((3, 3)))  # E: numpy.ndarray[numpy.float64]
+reveal_type(np.ones((3, 3), dtype=np.int64))  # E: numpy.ndarray[numpy.int64*]
+reveal_type(np.ones((3, 3), order='F'))  # E: numpy.ndarray[numpy.float64]

tests/reveal/ndarray_creation_py3.py

@@ -0,0 +1,4 @@
+import numpy as np
+
+nd: 'np.ndarray[np.int64]' = np.array([[1, 2], [3, 4]])
+reveal_type(nd)  # E: numpy.ndarray[numpy.int64]

tests/reveal/ndarray_methods.py

@@ -0,0 +1,4 @@
+import numpy as np
+
+nd = np.array([[1, 2], [3, 4]], dtype=np.int64)
+reveal_type(nd.dtype)  # E: numpy.int64