Phant
Phant

Reputation: 43

How to check input arguments in a python script with CLI?

I'm writing a small script to learn Python. The script prints a chess tournament table for N players. It has a simple CLI with a single argument N. Now I'm trying the following approach:

import argparse

def parse_args(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Tournament tables")
    parser.add_argument('N', help="number of players (2 at least)", type=int)
    args = parser.parse_args(argv)
    if args.N < 2:
        parser.error("N must be 2 at least")
    return args.N

def main(n: int) -> None:
    print(F"Here will be the table for {n} players")

if __name__ == '__main__':
    main(parse_args())

But this seems to have a flaw. The function main doesn't check n for invalid input (as it's the job of CLI parser). So if somebody calls main directly from another module (a tester for example), he may call it with lets say 0, and the program most likely crashes.

How should I properly handle this issue?

I'm considering several possible ways, but not sure what is the best.

  1. Add a proper value checking and error handling to main. This option looks ugly to me, as it violates the DRY principle and forces main to double the job of CLI.

  2. Just document that main must take only n >= 2, and its behaviour is unpredicted otherwise. Possibly to combine with adding an assertion check to main, like this:

    assert n >= 2, "n must be 2 or more"

  3. Perhaps such a function should not be external at all? So the whole chosen idiom is wrong and the script's entry point should be rewritten another way.

  4. ???

Upvotes: 4

Views: 1074

Answers (3)

JL Peyret
JL Peyret

Reputation: 12174

I’ve been using Pydantic liberally for enforcing data typing at runtime, within my code itself. Your N>=2 is easily enforced with a validator.

It’s a very robust, extremely widely used, library and very fast as it’s more of a data ingestion validator than a type checker.

And you could write the call as follows. How you call main is entirely up to you: direct call, argparse, click…


class ParamChecker(BaseModel):
    n : int

    @validator('n')
    def n2(cls, v):
        if v <2:
            raise ValueError('must be 2+')
        return v

def main(n: int) -> None:
    params = ParamChecker(**locals())

Pydantic also gets you informative, if not really end user friendly, error messages.

Upvotes: 1

Alex
Alex

Reputation: 7065

A common way of running argparse when wanting to test functions/CLI is to have the main function take a the sys.argv list and then call parse_args from within main like so:

arg.py

import argparse

def parse_args(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Tournament tables", prog="prog")
    parser.add_argument("N", help="number of players (2 at least)", type=int)
    args = parser.parse_args(argv)
    if args.N < 2:
        parser.error("N must be 2 at least")
    return args

def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    print(f"Here will be the table for {args.N} players")

if __name__ == "__main__":
    main()

This way a test can call main with a hypothetical CLI:

test_main.py

import pytest
from arg import main

def test_main(capsys):
    with pytest.raises(SystemExit):
        main(["0"])
    out, err = capsys.readouterr()
    assert err.splitlines()[-1] == "prog: error: N must be 2 at least"

Upvotes: 1

Thomas
Thomas

Reputation: 182000

You could have main do all the checking aind raise ArgumentError if something is amiss. Then catch that exception and forward it to the parser for display. Something along these lines:

import argparse

def run_with_args(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Tournament tables")
    parser.add_argument('N', help="number of players (2 at least)", type=int)
    args = parser.parse_args(argv)
    try:
        main(args.N)
    except argparse.ArgumentError as ex:
        parser.error(str(ex))

def main(n: int) -> None:
    if N < 2:
        raise argparse.ArgumentError("N must be 2 at least")
    print(F"Here will be the table for {n} players")

if __name__ == '__main__':
    run_with_args()

If you don't want to expose argparse.ArgumentError to library users of main, you can also create a custom exception type instead of it.

Upvotes: 3

Related Questions