From 3b4604a7d5e24a606ec0b65b3ad659a26fec1c48 Mon Sep 17 00:00:00 2001
From: Lars Reimann <mail@larsreimann.com>
Date: Sat, 22 Apr 2023 22:28:10 +0200
Subject: [PATCH] feat: `__getitem__` can accept a slice

---
 src/safeds/data/tabular/containers/_column.py | 15 ++++-
 .../data/tabular/exceptions/_exceptions.py    |  4 +-
 .../containers/_column/test_getitem.py        | 19 ------
 .../data/tabular/containers/test_column.py    | 59 +++++++++++++++++++
 4 files changed, 74 insertions(+), 23 deletions(-)
 delete mode 100644 tests/safeds/data/tabular/containers/_column/test_getitem.py

diff --git a/src/safeds/data/tabular/containers/_column.py b/src/safeds/data/tabular/containers/_column.py
index b12a9d976..f9ebff303 100644
--- a/src/safeds/data/tabular/containers/_column.py
+++ b/src/safeds/data/tabular/containers/_column.py
@@ -113,8 +113,19 @@ def __eq__(self, other: object) -> bool:
             return True
         return self.name == other.name and self._data.equals(other._data)
 
-    def __getitem__(self, index: int) -> _T:
-        return self.get_value(index)
+    def __getitem__(self, index: int | slice) -> _T:
+        if isinstance(index, int):
+            if index < 0 or index >= self._data.size:
+                raise IndexOutOfBoundsError(index)
+            return self._data[index]
+
+        if isinstance(index, slice):
+            if index.start < 0 or index.start > self._data.size:
+                raise IndexOutOfBoundsError(index)
+            if index.stop < 0 or index.stop > self._data.size:
+                raise IndexOutOfBoundsError(index)
+            data = self._data[index].reset_index(drop=True).rename(self.name)
+            return Column._from_pandas_series(data, self._type)
 
     def __iter__(self) -> Iterator[_T]:
         return iter(self._data)
diff --git a/src/safeds/data/tabular/exceptions/_exceptions.py b/src/safeds/data/tabular/exceptions/_exceptions.py
index 2a6a5cf8b..f8189ac19 100644
--- a/src/safeds/data/tabular/exceptions/_exceptions.py
+++ b/src/safeds/data/tabular/exceptions/_exceptions.py
@@ -42,11 +42,11 @@ class IndexOutOfBoundsError(IndexError):
 
     Parameters
     ----------
-    index : int
+    index : int | slice
         The wrongly used index.
     """
 
-    def __init__(self, index: int):
+    def __init__(self, index: int | slice):
         super().__init__(f"There is no element at index '{index}'.")
 
 
diff --git a/tests/safeds/data/tabular/containers/_column/test_getitem.py b/tests/safeds/data/tabular/containers/_column/test_getitem.py
deleted file mode 100644
index 2ae9d3a2b..000000000
--- a/tests/safeds/data/tabular/containers/_column/test_getitem.py
+++ /dev/null
@@ -1,19 +0,0 @@
-import pytest
-from safeds.data.tabular.containers import Column
-from safeds.data.tabular.exceptions import IndexOutOfBoundsError
-
-
-def test_getitem_valid() -> None:
-    column = Column("testColumn", [0, "1"])
-    assert column[0] == 0
-    assert column[1] == "1"
-
-
-# noinspection PyStatementEffect
-def test_getitem_invalid() -> None:
-    column = Column("testColumn", [0, "1"])
-    with pytest.raises(IndexOutOfBoundsError):
-        column[-1]
-
-    with pytest.raises(IndexOutOfBoundsError):
-        column[2]
diff --git a/tests/safeds/data/tabular/containers/test_column.py b/tests/safeds/data/tabular/containers/test_column.py
index 7da02851c..9932b324f 100644
--- a/tests/safeds/data/tabular/containers/test_column.py
+++ b/tests/safeds/data/tabular/containers/test_column.py
@@ -5,6 +5,7 @@
 import regex as re
 
 from safeds.data.tabular.containers import Column
+from safeds.data.tabular.exceptions import IndexOutOfBoundsError
 from safeds.data.tabular.typing import ColumnType, Integer, RealNumber, String, Boolean
 
 
@@ -68,6 +69,64 @@ def test_should_infer_type_if_not_passed(self, series: pd.Series, expected: Colu
         assert Column._from_pandas_series(series).type == expected
 
 
+class TestGetItem:
+    @pytest.mark.parametrize(
+        ("column", "index", "expected"),
+        [
+            (Column("a", [0, 1]), 0, 0),
+            (Column("a", [0, 1]), 1, 1),
+        ],
+        ids=[
+            "first item",
+            "second item"
+        ],
+    )
+    def test_should_get_the_item_at_index(self, column: Column, index: int, expected: Any) -> None:
+        assert column[index] == expected
+
+    @pytest.mark.parametrize(
+        ("column", "index", "expected"),
+        [
+            (Column("a", [0, 1, 2]), slice(0, 1), Column("a", [0])),
+            (Column("a", [0, 1, 2]), slice(2, 3), Column("a", [2])),
+            (Column("a", [0, 1, 2]), slice(0, 3), Column("a", [0, 1, 2])),
+            (Column("a", [0, 1, 2]), slice(0, 3, 2), Column("a", [0, 2])),
+        ],
+        ids=[
+            "first item",
+            "last item",
+            "all items",
+            "every other item"
+        ],
+    )
+    def test_should_get_the_items_at_slice(self, column: Column, index: slice, expected: Column) -> None:
+        assert column[index] == expected
+
+    @pytest.mark.parametrize(
+        "index",
+        [
+            -1,
+            2,
+            slice(-1, 2),
+            slice(0, 4),
+            slice(-1, 4)
+        ],
+        ids=[
+            "negative",
+            "out of bounds",
+            "slice with negative start",
+            "slice with out of bounds end",
+            "slice with negative start and out of bounds end"
+        ],
+    )
+    def test_should_raise_if_index_is_out_of_bounds(self, index: int | slice) -> None:
+        column = Column("a", [0, "1"])
+
+        with pytest.raises(IndexOutOfBoundsError):
+            # noinspection PyStatementEffect
+            column[index]
+
+
 class TestContains:
     @pytest.mark.parametrize(
         ("column", "value", "expected"),