From 075f96df2b3e361a73e2c8a08d96a801ff38adef Mon Sep 17 00:00:00 2001 From: Lars Reimann Date: Sat, 22 Apr 2023 22:37:30 +0200 Subject: [PATCH] fix: type hints --- src/safeds/data/tabular/containers/_column.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/safeds/data/tabular/containers/_column.py b/src/safeds/data/tabular/containers/_column.py index f9ebff303..c142273c3 100644 --- a/src/safeds/data/tabular/containers/_column.py +++ b/src/safeds/data/tabular/containers/_column.py @@ -3,7 +3,7 @@ import io from collections.abc import Sequence from numbers import Number -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar, overload import matplotlib.pyplot as plt import numpy as np @@ -82,7 +82,7 @@ def _from_pandas_series(data: pd.Series, type_: ColumnType | None = None) -> Col # Dunder methods # ------------------------------------------------------------------------------------------------------------------ - def __init__(self, name: str, data: Sequence) -> None: + def __init__(self, name: str, data: Sequence[_T]) -> None: """ Create a column. @@ -113,7 +113,15 @@ def __eq__(self, other: object) -> bool: return True return self.name == other.name and self._data.equals(other._data) - def __getitem__(self, index: int | slice) -> _T: + @overload + def __getitem__(self, index: int) -> _T: + ... + + @overload + def __getitem__(self, index: slice) -> Column[_T]: + ... + + def __getitem__(self, index: int | slice) -> _T | Column[_T]: if isinstance(index, int): if index < 0 or index >= self._data.size: raise IndexOutOfBoundsError(index)