From 6de9039315c8798f9d56f0045cb454767775138b Mon Sep 17 00:00:00 2001
From: Michael Hoang <enzime@users.noreply.github.com>
Date: Sun, 30 Mar 2025 23:11:03 +0900
Subject: [PATCH] lib/modules: add `class` to `specialArgs`

Co-Authored-By: Johannes Kirschbauer <hsjobeki@gmail.com>
---
 lib/modules.nix                               |  1 +
 lib/tests/modules.sh                          | 12 +++++
 .../modules/assert-module-class-is-nixos.nix  |  3 ++
 lib/tests/modules/class-check.nix             |  5 +-
 lib/tests/modules/expose-module-class.nix     |  8 +++
 lib/tests/modules/polymorphic-module.nix      | 23 ++++++++
 lib/tests/modules/specialArgs-class.nix       | 54 +++++++++++++++++++
 7 files changed, 105 insertions(+), 1 deletion(-)
 create mode 100644 lib/tests/modules/assert-module-class-is-nixos.nix
 create mode 100644 lib/tests/modules/expose-module-class.nix
 create mode 100644 lib/tests/modules/polymorphic-module.nix
 create mode 100644 lib/tests/modules/specialArgs-class.nix

diff --git a/lib/modules.nix b/lib/modules.nix
index 7716e855ebb2..290b9ce0b698 100644
--- a/lib/modules.nix
+++ b/lib/modules.nix
@@ -257,6 +257,7 @@ let
                     config
                     specialArgs
                     ;
+                  _class = class;
                 }
                 // specialArgs
               );
diff --git a/lib/tests/modules.sh b/lib/tests/modules.sh
index 9df6e61797b7..d0f25c283dcb 100755
--- a/lib/tests/modules.sh
+++ b/lib/tests/modules.sh
@@ -681,6 +681,18 @@ checkConfigOutput '^true$' config.viaConfig ./mkDefinition.nix
 checkConfigOutput '^true$' config.mkMerge ./mkDefinition.nix
 checkConfigOutput '^true$' config.mkForce ./mkDefinition.nix
 
+# specialArgs._class
+checkConfigOutput '"nixos"' config.nixos.config.foo ./specialArgs-class.nix
+checkConfigOutput '"bar"' config.conditionalImportAsNixos.config.foo ./specialArgs-class.nix
+checkConfigError 'attribute .*bar.* not found' config.conditionalImportAsNixos.config.bar ./specialArgs-class.nix
+checkConfigError 'attribute .*foo.* not found' config.conditionalImportAsDarwin.config.foo ./specialArgs-class.nix
+checkConfigOutput '"foo"' config.conditionalImportAsDarwin.config.bar ./specialArgs-class.nix
+checkConfigOutput '"nixos"' config.sub.nixos.foo ./specialArgs-class.nix
+checkConfigOutput '"bar"' config.sub.conditionalImportAsNixos.foo ./specialArgs-class.nix
+checkConfigError 'attribute .*bar.* not found' config.sub.conditionalImportAsNixos.bar ./specialArgs-class.nix
+checkConfigError 'attribute .*foo.* not found' config.sub.conditionalImportAsDarwin.foo ./specialArgs-class.nix
+checkConfigOutput '"foo"' config.sub.conditionalImportAsDarwin.bar ./specialArgs-class.nix
+
 cat <<EOF
 ====== module tests ======
 $pass Pass
diff --git a/lib/tests/modules/assert-module-class-is-nixos.nix b/lib/tests/modules/assert-module-class-is-nixos.nix
new file mode 100644
index 000000000000..9efd296b4135
--- /dev/null
+++ b/lib/tests/modules/assert-module-class-is-nixos.nix
@@ -0,0 +1,3 @@
+{ _class, ... }:
+assert _class == "nixos";
+{ }
diff --git a/lib/tests/modules/class-check.nix b/lib/tests/modules/class-check.nix
index 207f36315194..14c45263c96c 100644
--- a/lib/tests/modules/class-check.nix
+++ b/lib/tests/modules/class-check.nix
@@ -5,7 +5,9 @@
       nixosOk = lib.mkOption {
         type = lib.types.submoduleWith {
           class = "nixos";
-          modules = [ ];
+          modules = [
+            ./assert-module-class-is-nixos.nix
+          ];
         };
       };
       # Same but will have bad definition
@@ -45,6 +47,7 @@
       class = "nixos";
       modules = [
         ./module-class-is-nixos.nix
+        ./assert-module-class-is-nixos.nix
       ];
     };
 
diff --git a/lib/tests/modules/expose-module-class.nix b/lib/tests/modules/expose-module-class.nix
new file mode 100644
index 000000000000..f3b5147464d8
--- /dev/null
+++ b/lib/tests/modules/expose-module-class.nix
@@ -0,0 +1,8 @@
+{ _class, lib, ... }:
+{
+  options = {
+    foo = lib.mkOption {
+      default = _class;
+    };
+  };
+}
diff --git a/lib/tests/modules/polymorphic-module.nix b/lib/tests/modules/polymorphic-module.nix
new file mode 100644
index 000000000000..ad11002035d0
--- /dev/null
+++ b/lib/tests/modules/polymorphic-module.nix
@@ -0,0 +1,23 @@
+{ _class, lib, ... }:
+let
+  nixosModule =
+    { ... }:
+    {
+      options.foo = lib.mkOption {
+        default = "bar";
+      };
+    };
+  darwinModule =
+    { ... }:
+    {
+      options.bar = lib.mkOption {
+        default = "foo";
+      };
+    };
+in
+{
+  imports = [
+    (lib.optionalAttrs (_class == "nixos") nixosModule)
+    (lib.optionalAttrs (_class == "darwin") darwinModule)
+  ];
+}
diff --git a/lib/tests/modules/specialArgs-class.nix b/lib/tests/modules/specialArgs-class.nix
new file mode 100644
index 000000000000..847496472eed
--- /dev/null
+++ b/lib/tests/modules/specialArgs-class.nix
@@ -0,0 +1,54 @@
+{ lib, ... }:
+{
+  options = {
+    sub = {
+      nixos = lib.mkOption {
+        type = lib.types.submoduleWith {
+          class = "nixos";
+          modules = [
+            ./expose-module-class.nix
+          ];
+        };
+        default = { };
+      };
+
+      conditionalImportAsNixos = lib.mkOption {
+        type = lib.types.submoduleWith {
+          class = "nixos";
+          modules = [
+            ./polymorphic-module.nix
+          ];
+        };
+        default = { };
+      };
+
+      conditionalImportAsDarwin = lib.mkOption {
+        type = lib.types.submoduleWith {
+          class = "darwin";
+          modules = [
+            ./polymorphic-module.nix
+          ];
+        };
+        default = { };
+      };
+    };
+  };
+  config = {
+    _module.freeformType = lib.types.anything;
+
+    nixos = lib.evalModules {
+      class = "nixos";
+      modules = [ ./expose-module-class.nix ];
+    };
+
+    conditionalImportAsNixos = lib.evalModules {
+      class = "nixos";
+      modules = [ ./polymorphic-module.nix ];
+    };
+
+    conditionalImportAsDarwin = lib.evalModules {
+      class = "darwin";
+      modules = [ ./polymorphic-module.nix ];
+    };
+  };
+}