DataTable sort by function (or other callable)

CHANGELOG.md

@@ -16,8 +16,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - `MouseMove` events bubble up from widgets. `App` and `Screen` receive `MouseMove` events even if there's no Widget under the cursor. https://github.com/Textualize/textual/issues/2905
 
 ### Added
-- Added an interface for replacing prompt of an individual option in an `OptionList` https://github.com/Textualize/textual/issues/2603 
+- Added an interface for replacing prompt of an individual option in an `OptionList` https://github.com/Textualize/textual/issues/2603
 - Added `DirectoryTree.reload_node` method https://github.com/Textualize/textual/issues/2757
+- Added `key` argument to the `DataTable.sort()` method, allowing the table to be sorted using a custom function (or other callable). Providing both `columns` and a `key` will limit the data data sent to the key function to only those columns. https://github.com/Textualize/textual/issues/2261
 
 ## [0.32.0] - 2023-08-03
 

docs/examples/widgets/data_table_sort.py

@@ -0,0 +1,86 @@
+from rich.text import Text
+
+from textual.app import App, ComposeResult
+from textual.widgets import DataTable, Footer
+
+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"), 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),
+    (8, "Mehdy Metella", Text("France", style="italic"), 51.58, 52.15),
+    (7, "Tom Shields", Text("United States", style="italic"), 51.73, 51.12),
+    (1, "Aleksandr Sadovnikov", Text("Russia", style="italic"), 51.84, 50.85),
+    (10, "Darren Burns", Text("Scotland", style="italic"), 51.84, 51.55),
+]
+
+
+class TableApp(App):
+    BINDINGS = [
+        ("a", "sort_by_average_time", "Sort By Average Time"),
+        ("n", "sort_by_last_name", "Sort By Last Name"),
+        ("c", "sort_by_country", "Sort By Country"),
+    ]
+
+    current_sorts: set = set()
+
+    def compose(self) -> ComposeResult:
+        yield DataTable()
+        yield Footer()
+
+    def on_mount(self) -> None:
+        table = self.query_one(DataTable)
+        for col in ROWS[0]:
+            table.add_column(col, key=col)
+        table.add_rows(ROWS[1:])
+
+    def sort_reverse(self, sort_type: str):
+        """Determine if `sort_type` is ascending or descending."""
+        reverse = sort_type in self.current_sorts
+        if reverse:
+            self.current_sorts.remove(sort_type)
+        else:
+            self.current_sorts.add(sort_type)
+        return reverse
+
+    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_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_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(
+            "swimmer",
+            key=lambda swimmer: swimmer.split()[-1],
+            reverse=self.sort_reverse("swimmer"),
+        )
+
+    def action_sort_by_country(self) -> None:
+        """Sort DataTable by country which is a `Rich.Text` object."""
+        table = self.query_one(DataTable)
+        table.sort(
+            "country",
+            key=lambda country: country.plain,
+            reverse=self.sort_reverse("country"),
+        )
+
+
+app = TableApp()
+if __name__ == "__main__":
+    app.run()

docs/widgets/data_table.md

@@ -143,11 +143,22 @@ visible as you scroll through the data table.
 
 ### Sorting
 
-The `DataTable` can be sorted using the [sort][textual.widgets.DataTable.sort] method.
-In order to sort your data by a column, you must have supplied a `key` to the `add_column` method
-when you added it.
-You can then pass this key to the `sort` method to sort by that column.
-Additionally, you can sort by multiple columns by passing multiple keys to `sort`.
+The `DataTable` can be sorted using the [sort][textual.widgets.DataTable.sort] method. In order to sort your data by a column, you can provide the `key` you supplied to the `add_column` method or a `ColumnKey`. You can then pass one more column keys to the `sort` method to sort by one or more columns.
+
+Additionally, you can sort your `DataTable` with a custom function (or other callable) via the `key` argument. Similar to the `key` parameter of the built-in [sorted()](https://docs.python.org/3/library/functions.html#sorted) function, your function (or other callable) should take a single argument (row) and return a key to use for sorting purposes.
+
+Providing both `columns` and `key` will limit the row information sent to your `key` function (or other callable) to only the columns specified.
+
+=== "Output"
+
+    ```{.textual path="docs/examples/widgets/data_table_sort.py"}
+    ```
+
+=== "data_table_sort.py"
+
+    ```python
+    --8<-- "docs/examples/widgets/data_table_sort.py"
+    ```
 
 ### Labelled rows
 

src/textual/widgets/_data_table.py

@@ -4,7 +4,7 @@
 from dataclasses import dataclass
 from itertools import chain, zip_longest
 from operator import itemgetter
-from typing import Any, ClassVar, Generic, Iterable, NamedTuple, TypeVar, cast
+from typing import Any, Callable, ClassVar, Generic, Iterable, NamedTuple, TypeVar, cast
 
 import rich.repr
 from rich.console import RenderableType
@@ -2108,12 +2108,17 @@ def _get_fixed_offset(self) -> Spacing:
     def sort(
         self,
         *columns: ColumnKey | str,
+        key: Callable | None = None,
         reverse: bool = False,
     ) -> Self:
-        """Sort the rows in the `DataTable` by one or more column keys.
+        """Sort the rows in the `DataTable` by one or more column keys or a
+        key function (or other callable). If both columns and a key function
+        are specified, only data from those columns will sent to the key function.
 
         Args:
             columns: One or more columns to sort by the values in.
+            key: A function (or other callable) that returns a key to
+            use for sorting purposes.
             reverse: If True, the sort order will be reversed.
 
         Returns:
@@ -2127,11 +2132,19 @@ def sort_by_column_keys(
             result = itemgetter(*columns)(row_data)
             return result
 
+        _key = key
+        if key and columns:
+
+            def _key(row):
+                return key(itemgetter(*columns)(row[1]))
+
         ordered_rows = sorted(
-            self._data.items(), key=sort_by_column_keys, reverse=reverse
+            self._data.items(),
+            key=_key if _key else sort_by_column_keys,
+            reverse=reverse,
         )
         self._row_locations = TwoWayDict(
-            {key: new_index for new_index, (key, _) in enumerate(ordered_rows)}
+            {row_key: new_index for new_index, (row_key, _) in enumerate(ordered_rows)}
         )
         self._update_count += 1
         self.refresh()