diff --git a/datascience/tables.py b/datascience/tables.py index 243f5f751..1c0a84b20 100644 --- a/datascience/tables.py +++ b/datascience/tables.py @@ -148,7 +148,11 @@ def _add_column_and_format(self, table, label, column): ################# def __getitem__(self, label): - return self.values(label) + if isinstance(label, slice) or hasattr(label, '__index__'): + return self._with_columns(column[label] + for column in self._columns.values()) + else: + return self.values(label) def __setitem__(self, label, values): self.append_column(label, values) diff --git a/tests/test_tables.py b/tests/test_tables.py index a72bfe21d..20b7a9413 100644 --- a/tests/test_tables.py +++ b/tests/test_tables.py @@ -749,3 +749,13 @@ def test_q_chaining(table): letter | count | points a | 9 | 1 """) + +def test_row_slicing(table): + t = table[:2] + assert_equal(t.column_labels, ('letter', 'count', 'points')) + assert_array_equal(t.columns[0], ['a', 'b']) + assert_array_equal(t.columns[1], [9, 3]) + assert_array_equal(t.columns[2], [1, 2]) + + t_row = table[-1] + assert_equal(len(t_row.rows), 1)