From 3f914fc51f5ce9ea60c67f219c81c9cf00f2e3be Mon Sep 17 00:00:00 2001
From: Matthew Krausse <69082853+matthewkrausse@users.noreply.github.com>
Date: Thu, 29 Feb 2024 14:49:22 -0800
Subject: [PATCH] Add head and tail methods to Parsons Table (#1001)

* added head and tail methods to parsons table

* adding tests and docs to head and tail

---------

Co-authored-by: matthewkrausse <106627640+mattkrausse@users.noreply.github.com>
Co-authored-by: mkrausse-ggtx <131683556+mkrausse-ggtx@users.noreply.github.com>
---
 docs/table.rst     |  4 ++++
 parsons/etl/etl.py | 33 ++++++++++++++++++++++++++++++++-
 test/test_etl.py   | 12 ++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/docs/table.rst b/docs/table.rst
index 2c36e676bb..04a03c5cf9 100755
--- a/docs/table.rst
+++ b/docs/table.rst
@@ -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`
diff --git a/parsons/etl/etl.py b/parsons/etl/etl.py
index f2b20e94fe..6c91deeb4d 100644
--- a/parsons/etl/etl.py
+++ b/parsons/etl/etl.py
@@ -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
@@ -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
diff --git a/test/test_etl.py b/test/test_etl.py
index 89d8598aba..d6ed7ac8a4 100644
--- a/test/test_etl.py
+++ b/test/test_etl.py
@@ -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)