better example for docs

- Updated the example file for the docs to better show the functionality
of the change (especially when using `columns` and `key` together).
- Added one new tests to cover a similar situation to the example
  changes

docs/examples/widgets/data_table_sort.py

@@ -7,7 +7,7 @@
 ROWS = [
     ("lane", "swimmer", "country", "time 1", "time 2"),
     (4, "Joseph Schooling", Text("Singapore", style="italic"), 50.39, 51.84),
-    (2, "Michael Phelps", Text("United States", style="italic"), 51.14, 51.84),
+    (2, "Michael Phelps", Text("United States", style="italic"), 50.39, 51.84),
     (5, "Chad le Clos", Text("South Africa", style="italic"), 51.14, 51.73),
     (6, "László Cseh", Text("Hungary", style="italic"), 51.14, 51.58),
     (3, "Li Zhuhao", Text("China", style="italic"), 51.26, 51.26),
@@ -50,22 +50,25 @@ def action_sort_by_average_time(self) -> None:
         """Sort DataTable by average of times (via a function) and
         passing of column data through positional arguments."""
 
-        def sort_by_average_time(times):
-            return sum(n for n in times) / len(times)
+        def sort_by_average_time_then_last_name(row_data):
+            name, *scores = row_data
+            return (sum(scores) / len(scores), name.split()[-1])
 
         table = self.query_one(DataTable)
         table.sort(
+            "swimmer",
             "time 1",
             "time 2",
-            key=sort_by_average_time,
+            key=sort_by_average_time_then_last_name,
             reverse=self.sort_reverse("time"),
         )
 
     def action_sort_by_last_name(self) -> None:
         """Sort DataTable by last name of swimmer (via a lambda)."""
         table = self.query_one(DataTable)
         table.sort(
-            key=lambda row: row[1]["swimmer"].split()[-1],
+            "swimmer",
+            key=lambda swimmer: swimmer.split()[-1],
             reverse=self.sort_reverse("swimmer"),
         )
 

tests/test_data_table.py

@@ -1328,3 +1328,40 @@ def custom_sort(nums):
         )
         for i, row in enumerate(sorted_row_data):
             assert table.get_row_at(i) == row
+
+
+async def test_sort_by_multiple_columns_and_function():
+    """Test sorting a `DataTable` using a custom sort function and
+    only supplying data from specific columns."""
+
+    def custom_sort(row_data):
+        name, *scores = row_data
+        return (sum(scores) / len(scores), name.split()[-1])
+
+    row_data = (
+        ["ID", "Student", "Participation", "Test 1", "Test 2", "Test 3"],
+        ["ID-01", "Joseph Schooling", True, 90, 91, 92],
+        ["ID-02", "Li Zhuhao", False, 92, 93, 94],
+        ["ID-03", "Chad le Clos", False, 92, 93, 94],
+        ["ID-04", "Michael Phelps", True, 95, 96, 99],
+    )
+
+    app = DataTableApp()
+    async with app.run_test():
+        table = app.query_one(DataTable)
+
+        for col in row_data[0]:
+            table.add_column(col, key=col)
+        table.add_rows(row_data[1:])
+
+        table.sort("Student", "Test 1", "Test 2", "Test 3", key=custom_sort)
+        sorted_row_data = (row_data[1], row_data[3], row_data[2], row_data[4])
+        for i, row in enumerate(sorted_row_data):
+            assert table.get_row_at(i) == row
+
+        table.sort(
+            "Student", "Test 1", "Test 2", "Test 3", key=custom_sort, reverse=True
+        )
+        sorted_row_data = (row_data[4], row_data[2], row_data[3], row_data[1])
+        for i, row in enumerate(sorted_row_data):
+            assert table.get_row_at(i) == row