Repository: pri22296/beautifultable

Issue:
Allow sort based on an arbitrary callable
Currently we can sort the table based on a column. But we should allow the use of any callable as a key function like the sort method of list.

Comments:

Patch:
diff --git a/.gitignore b/.gitignore
index e70f354..87e4a82 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,4 @@ build/
 .tox/
 *.pyc
 .coverage
+.vscode/
diff --git a/beautifultable/beautifultable.py b/beautifultable/beautifultable.py
index dde020c..90ff9bd 100644
--- a/beautifultable/beautifultable.py
+++ b/beautifultable/beautifultable.py
@@ -800,12 +800,16 @@ class BeautifulTable(object):
             If `True` then table is sorted as if each comparison was reversed.
         """
         if isinstance(key, int):
-            index = key
+            key = operator.itemgetter(key)
         elif isinstance(key, basestring):
-            index = self.get_column_index(key)
+            key = operator.itemgetter(self.get_column_index(key))
+        elif callable(key):
+            pass
         else:
-            raise TypeError("'key' must either be 'int' or 'str'")
-        self._table.sort(key=operator.itemgetter(index), reverse=reverse)
+            raise TypeError(
+                "'key' must either be 'int' or 'str' or a 'callable'"
+            )
+        self._table.sort(key=key, reverse=reverse)
 
     def copy(self):
         """Return a shallow copy of the table.
@@ -1432,7 +1436,7 @@ class BeautifulTable(object):
 
         return "\n".join(string_)
 
-    def to_csv(self, file_name, delimiter=','):
+    def to_csv(self, file_name, delimiter=","):
         """Export table to CSV format.
 
         Parameters
@@ -1446,22 +1450,22 @@ class BeautifulTable(object):
 
         if not isinstance(file_name, str):
             raise ValueError(
-                (
-                    "Expected 'file_name' to be string, got {}"
-                ).format(type(file_name).__name__)
+                ("Expected 'file_name' to be string, got {}").format(
+                    type(file_name).__name__
+                )
             )
 
         try:
-            with open(file_name, mode='wt', newline='') as csv_file:
-                csv_writer = csv.writer(csv_file,
-                                        delimiter=delimiter,
-                                        quoting=csv.QUOTE_MINIMAL)
+            with open(file_name, mode="wt", newline="") as csv_file:
+                csv_writer = csv.writer(
+                    csv_file, delimiter=delimiter, quoting=csv.QUOTE_MINIMAL
+                )
                 csv_writer.writerow(self.column_headers)  # write header
                 csv_writer.writerows(self._table)  # write table
         except OSError:
             raise
 
-    def from_csv(self, file_name, delimiter=',', header_exists=True):
+    def from_csv(self, file_name, delimiter=",", header_exists=True):
         """Create table from CSV file.
 
         Parameters
@@ -1483,13 +1487,13 @@ class BeautifulTable(object):
 
         if not isinstance(file_name, str):
             raise ValueError(
-                (
-                    "Expected 'file_name' to be string, got {}"
-                ).format(type(file_name).__name__)
+                ("Expected 'file_name' to be string, got {}").format(
+                    type(file_name).__name__
+                )
             )
 
         try:
-            with open(file_name, mode='rt', newline='') as csv_file:
+            with open(file_name, mode="rt", newline="") as csv_file:
                 csv_file = csv.reader(csv_file, delimiter=delimiter)
 
                 if header_exists:


Test patch:
diff --git a/test.py b/test.py
index 0ad9017..2e8699f 100644
--- a/test.py
+++ b/test.py
@@ -25,6 +25,8 @@ class TableOperationsTestCase(unittest.TestCase):
         for item1, item2 in zip(iterable1, iterable2):
             self.assertEqual(item1, item2)
 
+    # Test for table operations
+
     def test_filter(self):
         new_table = self.table.filter(lambda x: x["rank"] > 1)
         self.assertEqual(len(self.table), 5)
@@ -36,6 +38,58 @@ class TableOperationsTestCase(unittest.TestCase):
         for row_t, row in zip(new_table, rows):
             self.compare_iterable(row_t, row)
 
+    def test_sort_by_index(self):
+        self.table.sort(0)
+        rows = [
+            ["Ethan", 2, "boy"],
+            ["Isabella", 1, "girl"],
+            ["Jacob", 1, "boy"],
+            ["Michael", 3, "boy"],
+            ["Sophia", 2, "girl"],
+        ]
+        for row_t, row in zip(self.table, rows):
+            self.compare_iterable(row_t, row)
+
+    def test_sort_by_index_reversed(self):
+        self.table.sort(0, reverse=True)
+        rows = [
+            ["Ethan", 2, "boy"],
+            ["Isabella", 1, "girl"],
+            ["Jacob", 1, "boy"],
+            ["Michael", 3, "boy"],
+            ["Sophia", 2, "girl"],
+        ]
+        for row_t, row in zip(self.table, reversed(rows)):
+            self.compare_iterable(row_t, row)
+
+    def test_sort_by_header(self):
+        self.table.sort("name")
+        rows = [
+            ["Ethan", 2, "boy"],
+            ["Isabella", 1, "girl"],
+            ["Jacob", 1, "boy"],
+            ["Michael", 3, "boy"],
+            ["Sophia", 2, "girl"],
+        ]
+        for row_t, row in zip(self.table, rows):
+            self.compare_iterable(row_t, row)
+
+    def test_sort_by_callable(self):
+        self.table.sort(lambda x: (x[1], x[0]))
+        rows = [
+            ["Isabella", 1, "girl"],
+            ["Jacob", 1, "boy"],
+            ["Ethan", 2, "boy"],
+            ["Sophia", 2, "girl"],
+            ["Michael", 3, "boy"],
+        ]
+        for row_t, row in zip(self.table, rows):
+            self.compare_iterable(row_t, row)
+
+    def test_sort_raises_exception(self):
+        with self.assertRaises(TypeError):
+            self.table.sort(None)
+
     # Tests for column operations
 
     def test_column_count(self):
@@ -530,26 +584,26 @@ class TableOperationsTestCase(unittest.TestCase):
 
     def test_csv_export(self):
         # Create csv files in path.
-        self.table.to_csv('beautiful_table.csv')
-        self.table.to_csv('./docs/beautiful_table.csv')
+        self.table.to_csv("beautiful_table.csv")
+        self.table.to_csv("./docs/beautiful_table.csv")
 
         with self.assertRaises(ValueError):
             self.table.to_csv(1)
 
         # Check if csv files exist.
-        self.assertTrue(os.path.exists('beautiful_table.csv'))
-        self.assertTrue(os.path.exists('./docs/beautiful_table.csv'))
+        self.assertTrue(os.path.exists("beautiful_table.csv"))
+        self.assertTrue(os.path.exists("./docs/beautiful_table.csv"))
 
         # Teardown step.
-        os.remove('beautiful_table.csv')
-        os.remove('./docs/beautiful_table.csv')
+        os.remove("beautiful_table.csv")
+        os.remove("./docs/beautiful_table.csv")
 
     def test_csv_import(self):
         # Export table as CSV file and import it back.
-        self.table.to_csv('beautiful_table.csv')
+        self.table.to_csv("beautiful_table.csv")
 
         test_table = BeautifulTable()
-        test_table.from_csv('beautiful_table.csv')
+        test_table.from_csv("beautiful_table.csv")
 
         with self.assertRaises(ValueError):
             self.table.from_csv(1)
@@ -561,11 +615,11 @@ class TableOperationsTestCase(unittest.TestCase):
         #     self.assertEqual(self.table[index], test_table[index])
 
         test_table = BeautifulTable()
-        test_table.from_csv('beautiful_table.csv', header_exists=False)
+        test_table.from_csv("beautiful_table.csv", header_exists=False)
         self.assertEqual(len(self.table), len(test_table) - 1)
 
         # Teardown step.
-        os.remove('beautiful_table.csv')
+        os.remove("beautiful_table.csv")
 
 
 if __name__ == "__main__":