Add shape manipulation methods to ndarray (#23)

* Add shape manipulation methods to ndarray

* Add tests for shape manipulation methods

Author: sproshev

numpy-stubs/__init__.pyi

@@ -10,6 +10,7 @@ from typing import (
     List,
     Mapping,
     Optional,
+    overload,
     Sequence,
     Sized,
     SupportsAbs,
@@ -368,6 +369,24 @@ class ndarray(_ArrayOrScalarCommon, Iterable, Sized, Container):
     @strides.setter
     def strides(self, value: _ShapeLike): ...
 
+    # Shape manipulation
+    @overload
+    def reshape(self, shape: Sequence[int], *, order: str=...) -> ndarray: ...
+    @overload
+    def reshape(self, *shape: int, order: str=...) -> ndarray: ...
+    @overload
+    def resize(self, new_shape: Sequence[int], *, refcheck: bool=...) -> None: ...
+    @overload
+    def resize(self, *new_shape: int, refcheck: bool=...) -> None: ...
+    @overload
+    def transpose(self, axes: Sequence[int]) -> ndarray: ...
+    @overload
+    def transpose(self, *axes: int) -> ndarray: ...
+    def swapaxes(self, axis1: int, axis2: int) -> ndarray: ...
+    def flatten(self, order: str=...) -> ndarray: ...
+    def ravel(self, order: str=...) -> ndarray: ...
+    def squeeze(self, axis: Union[int, Tuple[int, ...]]=...) -> ndarray: ...
+
     # Many of these special methods are irrelevant currently, since protocols
     # aren't supported yet. That said, I'm adding them for completeness.
     # https://docs.python.org/3/reference/datamodel.html

tests/pass/ndarray_shape_manipulation.py

@@ -0,0 +1,42 @@
+import numpy as np
+
+nd = np.array([[1, 2], [3, 4]])
+
+# reshape
+nd.reshape()
+nd.reshape(4)
+nd.reshape(2, 2)
+nd.reshape((2, 2))
+
+nd.reshape((2, 2), order="C")
+nd.reshape(4, order="C")
+
+# resize
+nd.resize()
+nd.resize(4)
+nd.resize(2, 2)
+nd.resize((2, 2))
+
+nd.resize((2, 2), refcheck=True)
+nd.resize(4, refcheck=True)
+
+# transpose
+nd.transpose()
+nd.transpose(1, 0)
+nd.transpose((1, 0))
+
+# swapaxes
+nd.swapaxes(0, 1)
+
+# flatten
+nd.flatten()
+nd.flatten("C")
+
+# ravel
+nd.ravel()
+nd.ravel("C")
+
+# squeeze
+nd.squeeze()
+nd.squeeze(0)
+nd.squeeze((0, 2))

tests/reveal/ndarray_shape_manipulation.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+nd = np.array([[1, 2], [3, 4]])
+
+# reshape
+reveal_type(nd.reshape())  # E: numpy.ndarray
+reveal_type(nd.reshape(4))  # E: numpy.ndarray
+reveal_type(nd.reshape(2, 2))  # E: numpy.ndarray
+reveal_type(nd.reshape((2, 2)))  # E: numpy.ndarray
+
+reveal_type(nd.reshape((2, 2), order="C"))  # E: numpy.ndarray
+reveal_type(nd.reshape(4, order="C"))  # E: numpy.ndarray
+
+# resize does not return a value
+
+# transpose
+reveal_type(nd.transpose())  # E: numpy.ndarray
+reveal_type(nd.transpose(1, 0))  # E: numpy.ndarray
+reveal_type(nd.transpose((1, 0)))  # E: numpy.ndarray
+
+# swapaxes
+reveal_type(nd.swapaxes(0, 1))  # E: numpy.ndarray
+
+# flatten
+reveal_type(nd.flatten())  # E: numpy.ndarray
+reveal_type(nd.flatten("C"))  # E: numpy.ndarray
+
+# ravel
+reveal_type(nd.ravel())  # E: numpy.ndarray
+reveal_type(nd.ravel("C"))  # E: numpy.ndarray
+
+# squeeze
+reveal_type(nd.squeeze())  # E: numpy.ndarray
+reveal_type(nd.squeeze(0))  # E: numpy.ndarray
+reveal_type(nd.squeeze((0, 2)))  # E: numpy.ndarray