Add option to pass Python args via main entry point

Python args are specified as a single string because if they were
positional, there wouldn't be a way to distinguish them from NsJail's
args.

2 files changed, 24 insertions, 8 deletions

diff --git a/DEVELOPING.md b/DEVELOPING.md
index 0dbf4cc..ef758bc 100644
--- a/DEVELOPING.md
+++ b/DEVELOPING.md
@@ -85,7 +85,7 @@ pipenv run devsh -c 'echo hello'
 NsJail can be invoked in a more direct manner that does not require using a web server or its API. See `python -m snekbox --help`. Example usage:
 
 ```bash
-python -m snekbox 'print("hello world!")' --time_limit 0
+python -m snekbox 'print("hello world!")' --time_limit 0 --- -m timeit
 ```
 
 With this command, NsJail uses the same configuration normally used through the web API. It also has an alias, `pipenv run eval`.
diff --git a/snekbox/__main__.py b/snekbox/__main__.py
index 0d3be6f..f92a48c 100644
--- a/snekbox/__main__.py
+++ b/snekbox/__main__.py
@@ -5,22 +5,38 @@ from snekbox.nsjail import NsJail
 
 def parse_args() -> argparse.Namespace:
     """Parse the command-line arguments and return the populated namespace."""
-    parser = argparse.ArgumentParser(prog="snekbox", usage="%(prog)s code [nsjail_args ...]")
+    parser = argparse.ArgumentParser(
+        prog="snekbox",
+        usage="%(prog)s [-h] code [nsjail_args ...] [--- py_args ...]",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
     parser.add_argument("code", help="the Python code to evaluate")
-    parser.add_argument("nsjail_args", nargs="?", help="override configured NsJail options")
-
-    # nsjail_args is just a dummy for documentation purposes.
-    # Its actual value comes from all the unknown arguments.
+    parser.add_argument(
+        "nsjail_args", nargs="?", default=[], help="override configured NsJail options"
+    )
+    parser.add_argument(
+        "py_args", nargs="?", default=["-c"], help="arguments to pass to the Python process"
+    )
+
+    # nsjail_args and py_args are just dummies for documentation purposes.
+    # Their actual values comes from all the unknown arguments.
     # There doesn't seem to be a better solution with argparse.
     args, unknown = parser.parse_known_args()
-    args.nsjail_args = unknown
+    try:
+        # Can't use double dash because that has special semantics for argparse already.
+        split = unknown.index("---")
+        args.nsjail_args = unknown[:split]
+        args.py_args = unknown[split + 1:]
+    except ValueError:
+        args.nsjail_args = unknown
+
     return args
 
 
 def main() -> None:
     """Evaluate Python code through NsJail."""
     args = parse_args()
-    result = NsJail().python3(args.code, nsjail_args=args.nsjail_args)
+    result = NsJail().python3(args.code, nsjail_args=args.nsjail_args, py_args=args.py_args)
     print(result.stdout)