Skip to content

Commit

Permalink
Add head and tail methods to Parsons Table (move-coop#1001)
Browse files Browse the repository at this point in the history
* added head and tail methods to parsons table

* adding tests and docs to head and tail

---------

Co-authored-by: matthewkrausse <[email protected]>
Co-authored-by: mkrausse-ggtx <[email protected]>
  • Loading branch information
3 people authored Feb 29, 2024
1 parent bb66a22 commit 3f914fc
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 1 deletion.
4 changes: 4 additions & 0 deletions docs/table.rst
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,10 @@ of commonly used methods. The full list can be found in the API section.

* - Method
- Description
* - :py:meth:`~parsons.etl.etl.ETL.head`
- Get the first n rows of a table
* - :py:meth:`~parsons.etl.etl.ETL.tail`
- Get the last n rows of a table
* - :py:meth:`~parsons.etl.etl.ETL.add_column`
- Add a column
* - :py:meth:`~parsons.etl.etl.ETL.remove_column`
Expand Down
33 changes: 32 additions & 1 deletion parsons/etl/etl.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,37 @@ class ETL(object):
def __init__(self):
pass

def head(self, n=5):
"""
Return the first n rows of the table
`Args:`
n: int
The number of rows to return. Defaults to 5.
`Returns:`
`Parsons Table`
"""

self.table = petl.head(self.table, n)

return self

def tail(self, n=5):
"""
Return the last n rows of the table. Defaults to 5.
`Args:`
n: int
The number of rows to return
`Returns:`
`Parsons Table`
"""

self.table = petl.tail(self.table, n)

return self

def add_column(self, column, value=None, index=None, if_exists="fail"):
"""
Add a column to your table
Expand Down Expand Up @@ -724,7 +755,7 @@ def long_table(
The new long table
"""

if type(key) == str:
if type(key) is str:
key = [key]

lt = self.cut(*key, column) # Create a table of key and column
Expand Down
12 changes: 12 additions & 0 deletions test/test_etl.py
Original file line number Diff line number Diff line change
Expand Up @@ -1000,3 +1000,15 @@ def test_deduplicate(self):
tbl_expected = Table([["a", "b", "c"], [1, 2, 3], [1, 3, 2], [2, 3, 4]])
tbl.deduplicate(["a", "b"])
assert_matching_tables(tbl_expected, tbl)

def test_head(self):
tbl = Table([["a", "b"], [1, 2], [3, 4], [5, 6], [7, 8], [9, 10]])
tbl_expected = Table([["a", "b"], [1, 2], [3, 4]])
tbl.head(2)
assert_matching_tables(tbl_expected, tbl)

def test_tail(self):
tbl = Table([["a", "b"], [1, 2], [3, 4], [5, 6], [7, 8], [9, 10]])
tbl_expected = Table([["a", "b"], [7, 8], [9, 10]])
tbl.tail(2)
assert_matching_tables(tbl_expected, tbl)

0 comments on commit 3f914fc

Please sign in to comment.