1 files changed, 105 insertions, 44 deletions

diff --git a/by-name-overlay.nix b/by-name-overlay.nix
index 3e52e64..0395492 100644
--- a/by-name-overlay.nix
+++ b/by-name-overlay.nix
@@ -12,22 +12,21 @@
 # This file should not depend on `pkgs` and thus not use `lib`.
 {
   baseDirectory,
-  fileName,
+  fileName ? null,
+  fileRegex ? null,
   finalizeFunction ? name: value: value,
   useShards ? true,
-}: let
-  # Takes a list of attrs as input and returns one merged attr set.
-  flattenAttrs = list:
-    if builtins.isList list
-    then
-      builtins.foldl' (acc: elem:
-        if builtins.isList elem
-        # Merging them with `//` is okay here, as we can be sure that the attr names are
-        # unique (they were separate dictionary after all)
-        then acc // (flattenAttrs elem)
-        else acc // elem) {}
-      list
-    else list;
+}:
+assert fileName == null -> fileRegex != null;
+assert fileRegex == null -> fileName != null; let
+  finalFileRegex =
+    if fileRegex == null
+    then "^${escapeRegex fileName}$"
+    else fileRegex;
+  fileDisplay =
+    if fileName == null
+    then "matching ${fileRegex}"
+    else fileName;
 
   # From nixpkgs/lib {{{
   # These functions are taken straight out of the `nixpkgs/lib`.
@@ -36,6 +35,12 @@
   mapAttrsToList = f: attrs:
     builtins.map (name: f name attrs.${name}) (builtins.attrNames attrs);
 
+  stringToCharacters = s: builtins.genList (p: builtins.substring p 1 s) (builtins.stringLength s);
+  escape = list: builtins.replaceStrings list (map (c: "\\${c}") list);
+  escapeRegex = escape (stringToCharacters "\\[{()^$?*+|.");
+
+  mapAttrs' = f: set: builtins.listToAttrs (map (attr: f attr set.${attr}) (builtins.attrNames set));
+
   nameValuePair = name: value: {inherit name value;};
   filterAttrs = pred: set:
     builtins.listToAttrs (builtins.concatMap (name: let
@@ -46,39 +51,95 @@
       else []) (builtins.attrNames set));
   # }}}
 
-  # Module files for a single shard
-  # Type: String -> String -> ListOf Path
-  namesForShard = shard: type:
-    if type != "directory"
-    then warn "Ignored non-directory, whilst importing by-name directory (${fileName}): '${shard}'" {}
-    else if useShards
-    then namesForElement shard type
-    else namesForElement "." type;
+  rawByName = {
+    baseDirectory,
+    finalFileRegex,
+    fileDisplay,
+    finalizeFunction,
+    useShards,
+  }: let
+    # Takes a list of attrs as input and returns one merged attr set.
+    flattenAttrs = list:
+      if builtins.isList list
+      then
+        builtins.foldl' (acc: elem:
+          if builtins.isList elem
+          # Merging them with `//` is okay here, as we can be sure that the attr names are
+          # unique (they were separate dictionary after all)
+          then acc // (flattenAttrs elem)
+          else acc // elem) {}
+        list
+      else list;
 
-  namesForElement = shard: _type: let
-    mkPath = name: type: let
-      path = baseDirectory + "/${shard}/${name}" + "/${fileName}";
+    # Module files for a single shard
+    # Type: String -> String -> ListOf Path
+    namesForShard = shard: type:
+      if type != "directory"
+      then warn "Ignored non-directory, whilst importing by-name directory (${fileDisplay}): '${shard}'" {}
+      else if useShards
+      then namesForElementShard shard type
+      else namesForElementDirect shard type;
+
+    # Type: String -> String -> ListOf Path
+    mkPath = name: toplevelType: let
+      rawPath = baseDirectory + "/${name}";
+      paths = filterAttrs (_: v: v != null) (mapAttrs' (name: value:
+        if builtins.match finalFileRegex name != null
+        then nameValuePair name value
+        else nameValuePair name null) (builtins.readDir rawPath));
+
+      checkPath = path: type:
+        if builtins.pathExists "${rawPath}/${path}"
+        then
+          if toplevelType != "directory"
+          then
+            # The `namesForShard` function should have already printed a warning.
+            [null]
+          else path
+        else warn "'${builtins.toString path}' does not exist. Skipped" null;
     in
-      if builtins.pathExists path
+      if toplevelType != "directory"
       then
-        if type != "directory"
-        then
-          # The `namesForShard` function should have already printed a warning.
-          null
-        else path
-      else warn "'${builtins.toString path}' does not exist. Skipped" null;
+        # The `namesForShard` function should have already printed a warning.
+        [null]
+      else if (builtins.attrValues paths) == []
+      then warn "'${fileDisplay}' did not match anything in ${rawPath}. Skipped" [null]
+      else mapAttrsToList checkPath paths;
+
+    # Type: String -> String -> AttrSet
+    namesForElementShard = shard: _type:
+      filterAttrs (name: value: value != null)
+      (builtins.mapAttrs
+        mkPath
+        (builtins.readDir (baseDirectory + "/${shard}")));
+
+    # Type: String -> String -> AttrSet
+    namesForElementDirect = name: type: {"${name}" = builtins.filter (value: value != null) (mkPath name type);};
+
+    # A list of all paths.
+    files = flattenAttrs (mapAttrsToList namesForShard (builtins.readDir baseDirectory));
+    output =
+      builtins.mapAttrs
+      finalizeFunction
+      files;
   in
-    filterAttrs (name: value: value != null)
-    (builtins.mapAttrs
-      mkPath
-      (builtins.readDir (baseDirectory + "/${shard}")));
+    output;
+
+  firstPass = rawByName {
+    inherit
+      baseDirectory
+      finalFileRegex
+      fileDisplay
+      finalizeFunction
+      useShards
+      ;
+  };
 
-  # A list of all module paths.
-  # These can the be simply injected into `import`
-  files = flattenAttrs (mapAttrsToList namesForShard (builtins.readDir baseDirectory));
-  output =
-    builtins.mapAttrs
-    finalizeFunction
-    files;
+  secondPass =
+    if fileName != null
+    then
+      builtins.mapAttrs (name: value: builtins.head value)
+      (filterAttrs (_: value: value != []) firstPass)
+    else firstPass;
 in
-  output
+  secondPass