From 11fb93763f175d104afe7737b62e765ad2931c6b Mon Sep 17 00:00:00 2001 From: Jayaram Kancherla Date: Fri, 26 Jan 2024 10:43:30 -0800 Subject: [PATCH] Improve `get_column` accessor (#97) --- src/biocframe/BiocFrame.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/biocframe/BiocFrame.py b/src/biocframe/BiocFrame.py index 232417a..e1ec4f1 100644 --- a/src/biocframe/BiocFrame.py +++ b/src/biocframe/BiocFrame.py @@ -620,10 +620,22 @@ def get_column(self, column: Union[str, int]) -> Any: The contents of the specified column. """ if isinstance(column, int): - column = self._column_names[column] - elif not isinstance(column, str): - raise TypeError("`column` must be either an integer index or column name.") - return self._data[column] + if column < 0: + raise IndexError("Index cannot be negative.") + + if column > len(self._column_names): + raise IndexError("Index greater than the number of columns.") + + return self._data[self._column_names[column]] + elif isinstance(column, str): + if column not in self._column_names: + raise AttributeError(f"Column: {column} does not exist.") + + return self._data[column] + + raise TypeError( + f"'column' must be a string or integer, provided '{type(column)}'." + ) def column(self, column: Union[str, int]) -> Any: """Alias for :py:meth:`~get_column`, provided for back-compatibility only."""