From 461f4a8708eed4ea6cf14ac6f775a197ec61eee8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marvin=20L=C3=B6bel?= <loebel.marvin@gmail.com>
Date: Tue, 2 Mar 2021 12:09:23 +0100
Subject: api: add convenience sort routines

This commit adds sort_by_key and sort_by_file_name convenience routines
for common sorting needs. The signature for sort_by_key in particular
mirrors the eponymous method on slices in the standard library.

Closes #147
---
 src/error.rs           |  2 +-
 src/lib.rs             | 40 ++++++++++++++++++++++++++++++++++++++--
 src/tests/recursive.rs | 45 ++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 3fb619c..9e25a07 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -98,7 +98,7 @@ impl Error {
     ///
     /// # Example
     ///
-    /// ```rust,no-run
+    /// ```rust,no_run
     /// use std::io;
     /// use std::path::Path;
     ///
diff --git a/src/lib.rs b/src/lib.rs
index f3c1da0..929c565 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -377,13 +377,14 @@ impl WalkDir {
         self
     }
 
-    /// Set a function for sorting directory entries.
+    /// Set a function for sorting directory entries with a comparator
+    /// function.
     ///
     /// If a compare function is set, the resulting iterator will return all
     /// paths in sorted order. The compare function will be called to compare
     /// entries from the same directory.
     ///
-    /// ```rust,no-run
+    /// ```rust,no_run
     /// use std::cmp;
     /// use std::ffi::OsString;
     /// use walkdir::WalkDir;
@@ -398,6 +399,41 @@ impl WalkDir {
         self
     }
 
+    /// Set a function for sorting directory entries with a key extraction
+    /// function.
+    ///
+    /// If a compare function is set, the resulting iterator will return all
+    /// paths in sorted order. The compare function will be called to compare
+    /// entries from the same directory.
+    ///
+    /// ```rust,no_run
+    /// use std::cmp;
+    /// use std::ffi::OsString;
+    /// use walkdir::WalkDir;
+    ///
+    /// WalkDir::new("foo").sort_by_key(|a| a.file_name().to_owned());
+    /// ```
+    pub fn sort_by_key<K, F>(self, mut cmp: F) -> Self
+    where
+        F: FnMut(&DirEntry) -> K + Send + Sync + 'static,
+        K: Ord,
+    {
+        self.sort_by(move |a, b| cmp(a).cmp(&cmp(b)))
+    }
+
+    /// Sort directory entries by file name, to ensure a deterministic order.
+    ///
+    /// This is a convenience function for calling `Self::sort_by()`.
+    ///
+    /// ```rust,no_run
+    /// use walkdir::WalkDir;
+    ///
+    /// WalkDir::new("foo").sort_by_file_name();
+    /// ```
+    pub fn sort_by_file_name(self) -> Self {
+        self.sort_by(|a, b| a.file_name().cmp(b.file_name()))
+    }
+
     /// Yield a directory's contents before the directory itself. By default,
     /// this is disabled.
     ///
diff --git a/src/tests/recursive.rs b/src/tests/recursive.rs
index bbb1ce1..4119f46 100644
--- a/src/tests/recursive.rs
+++ b/src/tests/recursive.rs
@@ -879,7 +879,7 @@ fn filter_entry() {
 }
 
 #[test]
-fn sort() {
+fn sort_by() {
     let dir = Dir::tmp();
     dir.mkdirp("foo/bar/baz/abc");
     dir.mkdirp("quux");
@@ -900,6 +900,49 @@ fn sort() {
     assert_eq!(expected, r.paths());
 }
 
+#[test]
+fn sort_by_key() {
+    let dir = Dir::tmp();
+    dir.mkdirp("foo/bar/baz/abc");
+    dir.mkdirp("quux");
+
+    let wd =
+        WalkDir::new(dir.path()).sort_by_key(|a| a.file_name().to_owned());
+    let r = dir.run_recursive(wd);
+    r.assert_no_errors();
+
+    let expected = vec![
+        dir.path().to_path_buf(),
+        dir.join("foo"),
+        dir.join("foo").join("bar"),
+        dir.join("foo").join("bar").join("baz"),
+        dir.join("foo").join("bar").join("baz").join("abc"),
+        dir.join("quux"),
+    ];
+    assert_eq!(expected, r.paths());
+}
+
+#[test]
+fn sort_by_file_name() {
+    let dir = Dir::tmp();
+    dir.mkdirp("foo/bar/baz/abc");
+    dir.mkdirp("quux");
+
+    let wd = WalkDir::new(dir.path()).sort_by_file_name();
+    let r = dir.run_recursive(wd);
+    r.assert_no_errors();
+
+    let expected = vec![
+        dir.path().to_path_buf(),
+        dir.join("foo"),
+        dir.join("foo").join("bar"),
+        dir.join("foo").join("bar").join("baz"),
+        dir.join("foo").join("bar").join("baz").join("abc"),
+        dir.join("quux"),
+    ];
+    assert_eq!(expected, r.paths());
+}
+
 #[test]
 fn sort_max_open() {
     let dir = Dir::tmp();
-- 
cgit v1.2.3