cli: Allow importing modules from the current directory when running a command or a file.

pyproject.toml

@@ -103,7 +103,7 @@ content-type = "text/markdown"
 file = "README.md"
 
 [project.scripts]
-sage = "sage.cli:main"
+sage = "sage.cli.__main__:main"
 
 [tool.conda-lock]
 platforms = ['linux-64', 'linux-aarch64', 'osx-64', 'osx-arm64']

src/sage/cli/__main__.py

@@ -1,5 +1,9 @@
+import os
 import sys
 
+# Allow importing modules from the current directory, matching python behavior
+sys.path.append(os.getcwd())
+
 from sage.cli import main
 
 sys.exit(main())

src/sage/cli/options.py

@@ -23,4 +23,4 @@ class CliOptions:
     command: str | None = None
 
     """The file to execute."""
-    file: str | None = None
+    file: list[str] | None = None

src/sage/cli/run_file_cmd.py

@@ -1,4 +1,5 @@
 import argparse
+import sys
 
 from sage.cli.options import CliOptions
 from sage.repl.preparse import preparse_file_named
@@ -23,7 +24,7 @@ def extend_parser(parser: argparse.ArgumentParser):
         """
         parser.add_argument(
             "file",
-            nargs="?",
+            nargs="*",
             help="execute the given file as sage code",
         )
 
@@ -32,13 +33,17 @@ def __init__(self, options: CliOptions):
         Initialize the command.
         """
         self.options = options
+        # shift sys.argv for compatibility with the old sage bash script and python command when consuming arguments from the command line
+        del sys.argv[0]
 
     def run(self) -> int:
         r"""
         Execute the given command.
         """
-        input_file = preparse_file_named(self.options.file) if self.options.file.endswith('.sage') else self.options.file
-        if self.options.file.endswith('.pyx') or self.options.file.endswith('.spyx'):
+        input_file = self.options.file[0]
+        if input_file.endswith('.sage'):
+            input_file = str(preparse_file_named(input_file))
+        if input_file.endswith('.pyx') or input_file.endswith('.spyx'):
             s = load_cython(input_file)
             eval(compile(s, tmp_filename(), 'exec'), sage_globals())
         else:

src/sage/cli/run_file_cmd_test.py

@@ -0,0 +1,27 @@
+from sage.cli.run_file_cmd import RunFileCmd
+from sage.cli.options import CliOptions
+from unittest.mock import patch
+import sys
+
+
+def test_run_file_cmd(capsys, tmp_path):
+    file = tmp_path / "test.sage"
+    file.write_text("print(3^33)")
+    options = CliOptions(file=[str(file)])
+    run_file_cmd = RunFileCmd(options)
+
+    result = run_file_cmd.run()
+    captured = capsys.readouterr()
+    assert captured.out == "5559060566555523\n"
+
+
+def test_run_file_cmd_with_args(capsys, tmp_path):
+    with patch.object(sys, 'argv', ["python3", "test.sage", "1", "1"]):
+        file = tmp_path / "test.sage"
+        file.write_text("import sys; print(int(sys.argv[1]) + int(sys.argv[2]))")
+        options = CliOptions(file=[str(file), "1", "1"])
+        run_file_cmd = RunFileCmd(options)
+
+        result = run_file_cmd.run()
+        captured = capsys.readouterr()
+        assert captured.out == "2\n"