Fix: Update dimension of a slice (#308)

src/skspatial/objects/_base_array.py

@@ -44,7 +44,7 @@ def __new__(cls: Type[Array], array: array_like) -> Array:
 
         return obj
 
-    def __array_wrap__(self, array, context=None):  # noqa: U100
+    def __array_wrap__(self, array: np.ndarray) -> np.ndarray:
         """
         Return regular :class:`numpy.ndarray` when default NumPy method is called.
 
@@ -60,37 +60,6 @@ def __array_wrap__(self, array, context=None):  # noqa: U100
         """
         return array
 
-    def __array_finalize__(self, obj: array_like) -> None:
-        """
-        Finalize creation of the array.
-
-        This function is required for adding extra attributes to a subclass of ndarray.
-        Without it, an array constructed from another may not have the extra attributes
-        (e.g., a projection of a vector onto another vector).
-
-        Examples
-        --------
-        >>> from skspatial.objects import Vector, Points
-
-        >>> vector_a = Vector([1, 0])
-        >>> vector_b = vector_a.project_vector([1, 1])
-
-        Without __array_finalize__, this vector will not have the dimension attribute.
-
-        >>> vector_a.dimension == vector_b.dimension
-        True
-
-        The same applies for 2D arrays.
-
-        >>> points = Points([[1, 2, 3], [4, 5, 6]])
-        >>> points_centered = points.mean_center()
-
-        >>> points.dimension == points_centered.dimension
-        True
-
-        """
-        self.dimension = getattr(obj, 'dimension', None)
-
     def to_array(self) -> np.ndarray:
         """
         Convert the object to a regular NumPy ndarray.
@@ -177,6 +146,10 @@ def __new__(cls: Type[Array1D], array: array_like) -> Array1D:
 
         return obj
 
+    def __array_finalize__(self, _) -> None:
+
+        self.dimension = self.size
+
     def set_dimension(self: Array1D, dim: int) -> Array1D:
         """
         Set the dimension (length) of the 1D array.
@@ -285,3 +258,11 @@ def set_dimension(self: Array2D, dim: int) -> Array2D:
         array_padded = np.pad(self, ((0, 0), (0, dim - self.dimension)), 'constant')
 
         return self.__class__(array_padded)
+
+    def __array_finalize__(self, _) -> None:
+
+        try:
+            self.dimension = self.shape[1]
+
+        except IndexError:
+            self.dimension = None

tests/unit/objects/test_base_array_1d.py

@@ -65,3 +65,14 @@ def test_dimension_failure(class_spatial, array, dimension):
 
     with pytest.raises(ValueError, match=message_expected):
         object_spatial.set_dimension(dimension)
+
+
+@pytest.mark.parametrize("class_spatial", [Point, Vector])
+def test_dimension_of_slice(class_spatial):
+
+    object_spatial = class_spatial([0, 0, 0])
+
+    assert object_spatial.dimension == 3
+    assert object_spatial[:3].dimension == 3
+    assert object_spatial[:2].dimension == 2
+    assert object_spatial[:1].dimension == 1

tests/unit/objects/test_points.py

@@ -52,6 +52,25 @@ def test_set_dimension(points, dim, points_expected):
     assert_array_equal(points.set_dimension(dim), points_expected)
 
 
+def test_dimension_of_slice():
+
+    points = Points([[0, 0, 0], [0, 0, 0]])
+
+    assert points.dimension == 3
+    assert points[:, :3].dimension == 3
+    assert points[:, :2].dimension == 2
+    assert points[:, :1].dimension == 1
+
+    assert points[:0, :].dimension == 3
+    assert points[:1, :].dimension == 3
+    assert points[:2, :].dimension == 3
+
+    # The dimension value is None here because the points are no longer a 2D array,
+    # so the normal method of finding the dimension fails.
+    assert points[:, 0].dimension is None
+    assert points[0, :].dimension is None
+
+
 @pytest.mark.parametrize(
     ("array_points", "array_centered_expected", "centroid_expected"),
     [