diff --git a/docs/tutorial.rst b/docs/tutorial.rst index 561b24c..117196d 100644 --- a/docs/tutorial.rst +++ b/docs/tutorial.rst @@ -146,6 +146,13 @@ To do so, we access the :class:`Dataset` as if it were a standard Python diction >>> data['First Name'] ['Kenneth', 'Bessie'] +You can also access the column using its index. :: + + >>> d.headers + ['Last Name', 'First Name', 'Age'] + >>> d.get_col(1) + ['Kenneth', 'Bessie'] + Let's find the average age. :: >>> ages = data['Age'] diff --git a/tablib/core.py b/tablib/core.py index 7e78b56..0429136 100644 --- a/tablib/core.py +++ b/tablib/core.py @@ -175,7 +175,6 @@ class Dataset(object): else: return [result.tuple for result in _results] - def __setitem__(self, key, value): self._validate(value) self._data[key] = Row(value) @@ -695,6 +694,12 @@ class Dataset(object): self.rpush_col(col, header) + def get_col(self, index): + """Returns the column from the :class:`Dataset` at the given index.""" + + return [row[index] for row in self._data] + + # ---- # Misc # ---- diff --git a/test_tablib.py b/test_tablib.py index 5fdf65a..48990a7 100755 --- a/test_tablib.py +++ b/test_tablib.py @@ -113,6 +113,22 @@ class TablibTestCase(unittest.TestCase): [self.john[2], self.george[2], self.tom[2]]) + def test_get_col(self): + """Verify getting columns by index""" + + self.assertEqual( + self.founders.get_col(self.headers.index('first_name')), + [self.john[0], self.george[0], self.tom[0]]) + + self.assertEqual( + self.founders.get_col(self.headers.index('last_name')), + [self.john[1], self.george[1], self.tom[1]]) + + self.assertEqual( + self.founders.get_col(self.headers.index('gpa')), + [self.john[2], self.george[2], self.tom[2]]) + + def test_data_slicing(self): """Verify slicing by data."""