lib.fileset.fileFilter: init

3 files changed, 144 insertions, 0 deletions

diff --git a/lib/fileset/default.nix b/lib/fileset/default.nix
index 7bd70167038..0342be3e037 100644
--- a/lib/fileset/default.nix
+++ b/lib/fileset/default.nix
@@ -6,6 +6,7 @@ let
     _coerceMany
     _toSourceFilter
     _unionMany
+    _fileFilter
     _printFileset
     _intersection
     ;
@@ -41,6 +42,7 @@ let
     ;
 
   inherit (lib.trivial)
+    isFunction
     pipe
     ;
 
@@ -279,6 +281,55 @@ If a directory does not recursively contain any file, it is omitted from the sto
       ];
 
   /*
+    Filter a file set to only contain files matching some predicate.
+
+    Type:
+      fileFilter ::
+        ({
+          name :: String,
+          type :: String,
+          ...
+        } -> Bool)
+        -> FileSet
+        -> FileSet
+
+    Example:
+      # Include all regular `default.nix` files in the current directory
+      fileFilter (file: file.name == "default.nix") ./.
+
+      # Include all non-Nix files from the current directory
+      fileFilter (file: ! hasSuffix ".nix" file.name) ./.
+
+      # Include all files that start with a "." in the current directory
+      fileFilter (file: hasPrefix "." file.name) ./.
+
+      # Include all regular files (not symlinks or others) in the current directory
+      fileFilter (file: file.type == "regular")
+  */
+  fileFilter =
+    /*
+      The predicate function to call on all files contained in given file set.
+      A file is included in the resulting file set if this function returns true for it.
+
+      This function is called with an attribute set containing these attributes:
+
+      - `name` (String): The name of the file
+
+      - `type` (String, one of `"regular"`, `"symlink"` or `"unknown"`): The type of the file.
+        This matches result of calling [`builtins.readFileType`](https://nixos.org/manual/nix/stable/language/builtins.html#builtins-readFileType) on the file's path.
+
+      Other attributes may be added in the future.
+    */
+    predicate:
+    # The file set to filter based on the predicate function
+    fileset:
+    if ! isFunction predicate then
+      throw "lib.fileset.fileFilter: Expected the first argument to be a function, but it's a ${typeOf predicate} instead."
+    else
+      _fileFilter predicate
+        (_coerce "lib.fileset.fileFilter: second argument" fileset);
+
+  /*
     The file set containing all files that are in both of two given file sets.
     See also [Intersection (set theory)](https://en.wikipedia.org/wiki/Intersection_(set_theory)).
 
diff --git a/lib/fileset/internal.nix b/lib/fileset/internal.nix
index 9892172955c..2d52a8cb410 100644
--- a/lib/fileset/internal.nix
+++ b/lib/fileset/internal.nix
@@ -638,4 +638,30 @@ rec {
     else
       # In all other cases it's the rhs
       rhs;
+
+  _fileFilter = predicate: fileset:
+    let
+      recurse = path: tree:
+        mapAttrs (name: subtree:
+          if isAttrs subtree || subtree == "directory" then
+            recurse (path + "/${name}") subtree
+          else if
+            predicate {
+              inherit name;
+              type = subtree;
+              # To ensure forwards compatibility with more arguments being added in the future,
+              # adding an attribute which can't be deconstructed :)
+              "lib.fileset.fileFilter: The predicate function passed as the first argument must be able to handle extra attributes for future compatibility. If you're using `{ name, file }:`, use `{ name, file, ... }:` instead." = null;
+            }
+          then
+            subtree
+          else
+            null
+        ) (_directoryEntries path tree);
+    in
+    if fileset._internalIsEmptyWithoutBase then
+      _emptyWithoutBase
+    else
+      _create fileset._internalBase
+        (recurse fileset._internalBase fileset._internalTree);
 }
diff --git a/lib/fileset/tests.sh b/lib/fileset/tests.sh
index 529f23ae887..d8d8dd41318 100755
--- a/lib/fileset/tests.sh
+++ b/lib/fileset/tests.sh
@@ -678,6 +678,73 @@ tree=(
 checkFileset 'intersection (unions [ ./a/b ./c/d ./c/e ]) (unions [ ./a ./c/d/f ./c/e ])'
 
 
+## File filter
+
+# The predicate is not called when there's no files
+tree=()
+checkFileset 'fileFilter (file: abort "this is not needed") ./.'
+checkFileset 'fileFilter (file: abort "this is not needed") _emptyWithoutBase'
+
+# The predicate must be able to handle extra attributes
+touch a
+expectFailure 'toSource { root = ./.; fileset = fileFilter ({ name, type }: true) ./.; }' 'called with unexpected argument '\''"lib.fileset.fileFilter: The predicate function passed as the first argument must be able to handle extra attributes for future compatibility. If you'\''re using `\{ name, file \}:`, use `\{ name, file, ... \}:` instead."'\'
+rm -rf -- *
+
+# .name is the name, and it works correctly, even recursively
+tree=(
+    [a]=1
+    [b]=0
+    [c/a]=1
+    [c/b]=0
+    [d/c/a]=1
+    [d/c/b]=0
+)
+checkFileset 'fileFilter (file: file.name == "a") ./.'
+tree=(
+    [a]=0
+    [b]=1
+    [c/a]=0
+    [c/b]=1
+    [d/c/a]=0
+    [d/c/b]=1
+)
+checkFileset 'fileFilter (file: file.name != "a") ./.'
+
+# `.type` is the file type
+mkdir d
+touch d/a
+ln -s d/b d/b
+mkfifo d/c
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type == "regular") ./.; }' \
+    'toSource { root = ./.; fileset = ./d/a; }'
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type == "symlink") ./.; }' \
+    'toSource { root = ./.; fileset = ./d/b; }'
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type == "unknown") ./.; }' \
+    'toSource { root = ./.; fileset = ./d/c; }'
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type != "regular") ./.; }' \
+    'toSource { root = ./.; fileset = union ./d/b ./d/c; }'
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type != "symlink") ./.; }' \
+    'toSource { root = ./.; fileset = union ./d/a ./d/c; }'
+expectEqual \
+    'toSource { root = ./.; fileset = fileFilter (file: file.type != "unknown") ./.; }' \
+    'toSource { root = ./.; fileset = union ./d/a ./d/b; }'
+rm -rf -- *
+
+# It's lazy
+tree=(
+    [b]=1
+    [c/a]=1
+)
+# Note that union evaluates the first argument first if necessary, that's why we can use ./c/a here
+checkFileset 'union ./c/a (fileFilter (file: assert file.name != "a"; true) ./.)'
+# but here we need to use ./c
+checkFileset 'union (fileFilter (file: assert file.name != "a"; true) ./.) ./c'
+
 ## Tracing
 
 # The second trace argument is returned