fix: support scripts with custom names (#1007)

Signed-off-by: Henry Schreiner <henryschreineriii@gmail.com>

Author: henryiii

nox/__main__.py

@@ -21,7 +21,7 @@
 
 from __future__ import annotations  # pragma: no cover
 
-from nox._cli import main  # pragma: no cover
+from nox._cli import nox_main as main  # pragma: no cover
 
 __all__ = ["main"]  # pragma: no cover
 

nox/_cli.py

@@ -31,14 +31,15 @@
 import nox.command
 import nox.virtualenv
 from nox import _options, tasks, workflow
+from nox._options import DefaultStr
 from nox._version import get_nox_version
 from nox.logger import logger, setup_logging
 from nox.project import load_toml
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
-__all__ = ["execute_workflow", "main"]
+__all__ = ["execute_workflow", "main", "nox_main"]
 
 
 def __dir__() -> list[str]:
@@ -136,7 +137,19 @@ def check_url_dependency(dep_url: str, dist: importlib.metadata.Distribution) ->
     return dep_purl.netloc == origin_purl.netloc and dep_purl.path == origin_purl.path
 
 
+def get_main_filename() -> str | None:
+    main_module = sys.modules.get("__main__")
+    if (
+        main_module
+        and (fname := getattr(main_module, "__file__", ""))
+        and os.path.exists(main_filename := os.path.abspath(fname))
+    ):
+        return main_filename
+    return None
+
+
 def run_script_mode(
+    noxfile: str,
     envdir: Path,
     *,
     reuse: bool,
@@ -163,11 +176,12 @@ def run_script_mode(
     subprocess.run([*cmd, *dependencies], env=env, check=True)
     nox_cmd = shutil.which("nox", path=env["PATH"])
     assert nox_cmd is not None, "Nox must be discoverable when installed"
+    args = [nox_cmd, "-f", noxfile, *sys.argv[1:]]
     # The os.exec functions don't work properly on Windows
     if sys.platform.startswith("win"):
         raise SystemExit(
             subprocess.run(
-                [nox_cmd, *sys.argv[1:]],
+                args,
                 env=env,
                 stdout=None,
                 stderr=None,
@@ -176,10 +190,18 @@ def run_script_mode(
                 check=False,
             ).returncode
         )
-    os.execle(nox_cmd, nox_cmd, *sys.argv[1:], env)  # pragma: nocover # noqa: S606
+    os.execle(nox_cmd, *args, env)  # pragma: nocover # noqa: S606
 
 
 def main() -> None:
+    _main(main_ep=False)
+
+
+def nox_main() -> None:
+    _main(main_ep=True)
+
+
+def _main(*, main_ep: bool) -> None:
     args = _options.options.parse_args()
 
     if args.help:
@@ -198,7 +220,12 @@ def main() -> None:
         msg = f"Invalid NOX_SCRIPT_MODE: {nox_script_mode!r}, must be one of 'none', 'reuse', or 'fresh'"
         raise SystemExit(msg)
     if nox_script_mode != "none":
-        toml_config = load_toml(os.path.expandvars(args.noxfile), missing_ok=True)
+        noxfile = (
+            args.noxfile
+            if main_ep or not isinstance(args.noxfile, DefaultStr)
+            else (get_main_filename() or args.noxfile)
+        )
+        toml_config = load_toml(os.path.expandvars(noxfile), missing_ok=True)
         dependencies = toml_config.get("dependencies")
         if dependencies is not None:
             valid_env = check_dependencies(dependencies)
@@ -235,6 +262,7 @@ def main() -> None:
 
                 envdir = Path(args.envdir or ".nox")
                 run_script_mode(
+                    noxfile,
                     envdir,
                     reuse=nox_script_mode == "reuse",
                     dependencies=dependencies,

nox/_options.py

@@ -42,6 +42,11 @@ def __dir__() -> list[str]:
     return __all__
 
 
+# User-specified arguments will be a regular string
+class DefaultStr(str):
+    __slots__ = ()
+
+
 ReuseVenvType = Literal["no", "yes", "never", "always"]
 
 options = _option_set.OptionSet(
@@ -515,7 +520,7 @@ def _tag_completer(
         "-f",
         "--noxfile",
         group=options.groups["general"],
-        default="noxfile.py",
+        default=DefaultStr("noxfile.py"),
         help="Location of the Python file containing Nox sessions.",
     ),
     _option_set.Option(

nox/tasks.py

@@ -68,7 +68,7 @@ def _load_and_exec_nox_module(global_config: Namespace) -> types.ModuleType:
         types.ModuleType: The initialised Nox module.
     """
     spec = importlib.util.spec_from_file_location(
-        "user_nox_module", global_config.noxfile
+        "user_nox_module", str(global_config.noxfile)
     )
     assert spec is not None  # If None, fatal importlib error, would crash anyway
 
@@ -107,8 +107,9 @@ def load_nox_module(global_config: Namespace) -> types.ModuleType | int:
     # Save the absolute path to the Noxfile.
     # This will inoculate it if Nox changes paths because of an implicit
     # or explicit chdir (like the one below).
-    global_config.noxfile = os.path.join(
-        noxfile_parent_dir, os.path.basename(global_config_noxfile)
+    # Keeps the class of the original string
+    global_config.noxfile = global_config.noxfile.__class__(
+        os.path.join(noxfile_parent_dir, os.path.basename(global_config_noxfile))
     )
 
     try:

tests/resources/noxfile_script_mode_exec.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+
+# /// script
+# dependencies = ["nox", "cowsay"]
+# ///
+
+
+import nox
+
+
+@nox.session
+def exec_example(session: nox.Session) -> None:
+    # Importing inside the function so that if the test fails,
+    # it shows a better failure than immediately failing to import
+    import cowsay  # noqa: PLC0415
+
+    print(cowsay.cow("another_world"))
+
+
+if __name__ == "__main__":
+    nox.main()

tests/test__cli.py

@@ -87,6 +87,8 @@ def test_invalid_backend_envvar(
 ) -> None:
     monkeypatch.setenv("NOX_SCRIPT_VENV_BACKEND", "invalid")
     monkeypatch.setattr(sys, "argv", ["nox"])
+    # This will return pytest's filename instead, so patching it to None
+    monkeypatch.setattr(nox._cli, "get_main_filename", lambda: None)
     monkeypatch.chdir(tmp_path)
     tmp_path.joinpath("noxfile.py").write_text(
         "# /// script\n# dependencies=['nox', 'invalid']\n# ///",
@@ -101,6 +103,8 @@ def test_invalid_backend_inline(
     monkeypatch: pytest.MonkeyPatch, tmp_path: Path
 ) -> None:
     monkeypatch.setattr(sys, "argv", ["nox"])
+    # This will return pytest's filename instead, so patching it to None
+    monkeypatch.setattr(nox._cli, "get_main_filename", lambda: None)
     monkeypatch.chdir(tmp_path)
     tmp_path.joinpath("noxfile.py").write_text(
         "# /// script\n# dependencies=['nox', 'invalid']\n# tool.nox.script-venv-backend = 'invalid'\n# ///",

tests/test_main.py

@@ -27,6 +27,7 @@
 import pytest
 
 import nox
+import nox._cli
 import nox._options
 import nox.registry
 import nox.sessions
@@ -44,14 +45,17 @@
 os.environ.pop("NOXSESSION", None)
 
 
-def test_main_no_args(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize(
+    "main", [nox.main, nox._cli.nox_main], ids=["main", "nox_main"]
+)
+def test_main_no_args(monkeypatch: pytest.MonkeyPatch, main: Any) -> None:
     monkeypatch.setattr(sys, "argv", [sys.executable])
     with mock.patch("nox.workflow.execute") as execute:
         execute.return_value = 0
 
         # Call the function.
         with mock.patch.object(sys, "exit") as exit:
-            nox.main()
+            main()
             exit.assert_called_once_with(0)
         assert execute.called
 
@@ -1098,6 +1102,26 @@ def test_noxfile_no_script_mode(monkeypatch: pytest.MonkeyPatch) -> None:
     assert "No module named 'cowsay'" in job.stderr
 
 
+def test_noxfile_script_mode_exec(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("NOX_SCRIPT_MODE", raising=False)
+    job = subprocess.run(
+        [
+            sys.executable,
+            Path(RESOURCES) / "noxfile_script_mode_exec.py",
+            "-s",
+            "exec_example",
+        ],
+        check=False,
+        capture_output=True,
+        text=True,
+        encoding="utf-8",
+    )
+    print(job.stdout)
+    print(job.stderr)
+    assert job.returncode == 0
+    assert "another_world" in job.stdout
+
+
 def test_noxfile_script_mode_url_req() -> None:
     job = subprocess.run(
         [