append_column now returns the table modified as well

data-8 · May 16, 2017 · d713ab3 · d713ab3
1 parent 7a7ca82
commit d713ab3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/datascience/tables.py b/datascience/tables.py
@@ -449,14 +449,12 @@ def append_column(self, label, values):
         c      | 3     | 2
         z      | 1     | 10
         >>> table.append_column('new_col1', make_array(10, 20, 30, 40))
-        >>> table
         letter | count | points | new_col1
         a      | 9     | 1      | 10
         b      | 3     | 2      | 20
         c      | 3     | 2      | 30
         z      | 1     | 10     | 40
         >>> table.append_column('new_col2', 'hello')
-        >>> table
         letter | count | points | new_col1 | new_col2
         a      | 9     | 1      | 10       | hello
         b      | 3     | 2      | 20       | hello
@@ -491,6 +489,7 @@ def append_column(self, label, values):
             self._num_rows = len(values)
 
         self._columns[label] = values
+        return self
 
     def relabel(self, column_label, new_label):
         """Changes the label(s) of column(s) specified by ``column_label`` to

diff --git a/tests/test_tables.py b/tests/test_tables.py
@@ -585,15 +585,21 @@ def test_append_column(table):
     c      | 3     | 2      | 30
     z      | 1     | 10     | 40
     """)
-    table.append_column('new_col2', column_2)
-    print(table)
+    ret_table = table.append_column('new_col2', column_2)
     assert_equal(table, """
     letter | count | points | new_col1 | new_col2
     a      | 9     | 1      | 10       | hello
     b      | 3     | 2      | 20       | hello
     c      | 3     | 2      | 30       | hello
     z      | 1     | 10     | 40       | hello
     """)
+    assert_equal(ret_table, """
+    letter | count | points | new_col1 | new_col2
+    a      | 9     | 1      | 10       | hello
+    b      | 3     | 2      | 20       | hello
+    c      | 3     | 2      | 30       | hello
+    z      | 1     | 10     | 40       | hello
+    """)
 
     with(pytest.raises(ValueError)):
         table.append_column('bad_col', [1, 2])