DEV Community

Cover image for The Ultimate Python main()
Rémy 🤖
Rémy 🤖

Posted on

The Ultimate Python main()

So you want to write a CLI utility in Python. The goal is to be able to write:

$ ./say_something.py -w "hello, world"
hello, world
Enter fullscreen mode Exit fullscreen mode

And then to get your script running. This isn't so hard and you've probably done it already, but have you thought of all the use cases? Let's go step by step on all the things you can do to make your main() bullet-proof. The first sections are going to be very basic and we'll build up something less obvious as we progress.

Hello, world

Let's start with a basic script that just says "hello, wold".

print("hello, world")
Enter fullscreen mode Exit fullscreen mode

Now this is fairly basic but still won't work if you call it directly like ./say_something.py. In order for this to work, you need to add the famous shebang at the beginning.

#!/usr/bin/python3

print("hello, world")
Enter fullscreen mode Exit fullscreen mode

Note — Let's also not forget to do a chmod a+x say_something.py to give execution rights on the file

However, what if you're using a virtualenv? By calling /usr/bin/python3 directly you're forcing the path to the Python interpreter. This can be desirable in some cases (by example package managers will tend to force this for safety) but if you're writing a distributable script it's simpler to get it from the environment.

#!/usr/bin/env python3

print("hello, world")
Enter fullscreen mode Exit fullscreen mode

Note — We're calling python3 instead of python because in some systems you'll find that python points to Python 2 and who wants to be using Python 2?

Importability

Something that you'll learn early on from manuals is that you want to make sure that if for some reason someone imports your module then there won't be any side-effect like a surprise print().

Because of that, you need to "protect" your code with this special idiom:

#!/usr/bin/env python3

if __name__ == "__main__":
    print("hello, world")
Enter fullscreen mode Exit fullscreen mode

You'll note the use of the special variable __name__ whose value is going to be __main__ if you're in the "root" file that Python is running and the module name otherwise.

Parsing arguments

You'll notice that so far the program only says "hello, world" but doesn't parse the -w argument that we've described above. Parsing arguments in Python is not so different from C — for those who had the chance of learning it at school — except that well it's Python so you have a few more goodies.

The main difference is that instead of receiving the arguments as a an argument of main() you'll have to import them.

#!/usr/bin/env python3
from sys import argv, stderr


def fail(msg: str):
    stderr.write(f"{msg}\n")
    exit(1)


if __name__ == "__main__":
    what = "hello, world"

    if len(argv) == 1:
        pass
    elif len(argv) == 3:
        if argv[1] != "-w":
            fail(f'Unrecognized argument "{argv[1]}"')
        what = argv[2]
    else:
        fail("Too many/few arguments")

    print(what)
Enter fullscreen mode Exit fullscreen mode

Now, while this fits the bill, this is oh so much more laborious to parse arguments manually like this. For a long time I was afraid to use argparse because the documentation is quite the beast, but in the end its basic use is very simple.

#!/usr/bin/env python3
from argparse import ArgumentParser


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    args = parser.parse_args()

    print(args.what)
Enter fullscreen mode Exit fullscreen mode

That way you let all the hard work to ArgumentParser() and as a bonus you get an auto-generated help text when calling with -h as an argument.

Outside call

All this is cool but what if you want to call the CLI tool from another Python program? Experience showed me that it's often a much more pragmatic thing to do rather than to design a specific API for Python and a different API for CLI (which is arguments parsing).

The problem with what we've got here is that we can't reference our call from the outside world. To run it has to be in __main__, which is not necessarily what we want.

Fortunately there is an easy solution to this: create a C-style main() function.

#!/usr/bin/env python3
from argparse import ArgumentParser


def main():
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    args = parser.parse_args()

    print(args.what)


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode

While this works, this doesn't fit the bill of letting you call custom arguments from an outside package. However, all you need to do is add the arguments as an argument to your main function:

#!/usr/bin/env python3
from argparse import ArgumentParser
from typing import Sequence, Optional


def main(argv: Optional[Sequence[str]] = None):
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    args = parser.parse_args(argv)

    print(args.what)


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode

The advantage of this is that by default it will keep on working like before except that this time it will let you pass custom arguments to call it from the outside, like that by example:

>>> from say_something import main
>>> main(["-w", "hello, world"])
hello, world
Enter fullscreen mode Exit fullscreen mode

A bit of assistance

For good measure and clarity, let's take out the arguments parsing into a separate function.

#!/usr/bin/env python3
from argparse import ArgumentParser, Namespace
from typing import Sequence, Optional


def parse_args(argv: Optional[Sequence[str]] = None) -> Namespace:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return parser.parse_args(argv)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)
    print(args.what)


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode

But now you'll notice that the typing annotation isn't very helpful. I usually like to be a little bit more verbose to get more assistance from my IDE later on.

#!/usr/bin/env python3
from argparse import ArgumentParser
from typing import Sequence, Optional, NamedTuple


class Args(NamedTuple):
    what: str


def parse_args(argv: Optional[Sequence[str]] = None) -> Args:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return Args(**parser.parse_args(argv).__dict__)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)
    print(args.what)


if __name__ == "__main__":
    main()
Enter fullscreen mode Exit fullscreen mode

Handling signals

For the purpose of this demonstration, let's add some "sleep" in this code to simulate that it's doing something instead of printing the text right away.

Now what if the program receives a signal? There is two main things you want to handle:

  • SIGINT — When the user does a CTRL+C in their terminal, which raises a KeyboardInterrupt
  • SIGTERM — When the user kindly asks the program to die with a TERM signal, which can be handled (as opposed to SIGKILL) but we'll see that later

Let's start by handling CTRL+C:

#!/usr/bin/env python3
from argparse import ArgumentParser
from time import sleep
from typing import Sequence, Optional, NamedTuple
from sys import stderr


class Args(NamedTuple):
    what: str


def parse_args(argv: Optional[Sequence[str]] = None) -> Args:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return Args(**parser.parse_args(argv).__dict__)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)
    sleep(100)
    print(args.what)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        stderr.write("ok, bye\n")
        exit(1)
Enter fullscreen mode Exit fullscreen mode

As you can see, it's as simple as catching the KeyboardInterrupt exception. Please note that we do this outside of main() because the famous hypothetical "caller module" will already have its signal handling in place so that's something that we only need to setup if we're running on ourself.

Next comes the handling of SIGTERM. By default the program will just stop but we don't really want this. By example if we're using the following pattern inside the code:

try:
    # do something
finally:
    # cleanup
Enter fullscreen mode Exit fullscreen mode

We won't have a chance to cleanup if an exception isn't raised. Fortunately for us, we can raise the exception for ourselves.

#!/usr/bin/env python3
from argparse import ArgumentParser
from time import sleep
from typing import Sequence, Optional, NamedTuple
from sys import stderr
from signal import signal, SIGTERM


class Args(NamedTuple):
    what: str


def parse_args(argv: Optional[Sequence[str]] = None) -> Args:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return Args(**parser.parse_args(argv).__dict__)


def sigterm_handler(_, __):
    raise SystemExit(1)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)
    sleep(100)
    print(args.what)


if __name__ == "__main__":
    signal(SIGTERM, sigterm_handler)

    try:
        main()
    except KeyboardInterrupt:
        stderr.write("ok, bye\n")
        exit(1)
Enter fullscreen mode Exit fullscreen mode

What happens is that we're using signal() to register sigterm_handler() as our handler for SIGTERM. What happens "under the hood" is that the signal handler will be executed before the "next instruction" that the interpreter would otherwise have considered. This gives you a chance to raise an exception which is going to bubble up to our __main__ and exit the program with a 1 return code while triggering all the finally and context managers along the way.

As stated before, there is a sleep() in the middle of the function. This means that you can run the code from this section and then hit CTLR+C or send a SIGTERM to see what happens when you interrupt the program.

Error Reporting

Sometimes — more often that you'd like — your program will fail either by its own fault or because the inputs are incorrect. By example, if the user tries to open a file that doesn't exist, you might want to report it. And by "report it" I mean nicely, not with a harsh stack trace. Keep the stack trace for the cases that you didn't expect this way you know that something is really wrong.

Let's imagine that we want to forbid the user to say "ni". In order to report the error, we're going to create a specific error type for our program

class SaySomethingError(Exception):
    pass
Enter fullscreen mode Exit fullscreen mode

Then we're going to handle it in our __main__:

    except SaySomethingError as e:
        stderr.write(f"Error: {e}")
        exit(1)
Enter fullscreen mode Exit fullscreen mode

And finally we're going to raise the error from main()

    if args.what == "ni":
        raise SaySomethingError('Saying "ni" is forbidden')
Enter fullscreen mode Exit fullscreen mode

For an overall code that is:

#!/usr/bin/env python3
from argparse import ArgumentParser
from typing import Sequence, Optional, NamedTuple
from sys import stderr
from signal import signal, SIGTERM


class Args(NamedTuple):
    what: str


class SaySomethingError(Exception):
    pass


def parse_args(argv: Optional[Sequence[str]] = None) -> Args:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return Args(**parser.parse_args(argv).__dict__)


def sigterm_handler(_, __):
    raise SystemExit(1)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)

    if args.what == "ni":
        raise SaySomethingError('Saying "ni" is forbidden')

    print(args.what)


if __name__ == "__main__":
    signal(SIGTERM, sigterm_handler)

    try:
        main()
    except KeyboardInterrupt:
        stderr.write("ok, bye\n")
        exit(1)
    except SaySomethingError as e:
        stderr.write(f"Error: {e}")
        exit(1)
Enter fullscreen mode Exit fullscreen mode

With this very simple system, you can raise an error from your code no matter how deep you are. It will trigger all the cleanup functions from context managers and finally blocks. And finally it will be caught by the __main__ to display a proper error message.

Packaging

A good idea when providing a bin script inside a package is to have it be called with Python's -m option. By example instead of writing pip, I usually write python3 -m pip to be sure that the Pip I'm running is indeed the one from my virtual env. As a bonus, you don't need the environment's bin directory in your $PATH. You'll find that most famous and not-so-famous packages provide both ways of calling their binaries.

In order to do that, you need to put your script into a __main__.py file. Let's do this (this is UNIX commands but Windows should have no trouble translating):

mkdir say_something
touch say_something/__init__.py
mv say_something.py say_something/__main__.py
Enter fullscreen mode Exit fullscreen mode

Now you can call your script with python3 -m say_something.

Note — We're creating an empty __init__.py file to signify to Python that this is a proper module

Here it's clearly becoming a question of preferences but my personal favorite tool for packaging has become Poetry because it is so simple. Let's create a basic pyproject.toml for our project.

[tool.poetry]
name = "say_something"
version = "0.1.0"
description = ""
authors = []
license = "WTFPL"

[tool.poetry.dependencies]
python = "^3.8"

[tool.poetry.dev-dependencies]

[build-system]
requires = ["poetry>=0.12"]
build-backend = "poetry.masonry.api"
Enter fullscreen mode Exit fullscreen mode

Let's try this out to see if it worked

poetry run python -m say_something
Enter fullscreen mode Exit fullscreen mode

This is good but did the command go? Have no fear, Poetry lets you declare your "bin" commands pretty easily. The only thing is that it needs a function to call directly, including all the signal-handling shenanigans. Let's move it all into a separate function then:

#!/usr/bin/env python3
from argparse import ArgumentParser
from typing import Sequence, Optional, NamedTuple
from sys import stderr
from signal import signal, SIGTERM


class Args(NamedTuple):
    what: str


class SaySomethingError(Exception):
    pass


def parse_args(argv: Optional[Sequence[str]] = None) -> Args:
    parser = ArgumentParser()
    parser.add_argument("-w", "--what", default="hello, world")

    return Args(**parser.parse_args(argv).__dict__)


def sigterm_handler(_, __):
    raise SystemExit(1)


def main(argv: Optional[Sequence[str]] = None):
    args = parse_args(argv)

    if args.what == "ni":
        raise SaySomethingError('Saying "ni" is forbidden')

    print(args.what)


def __main__():
    signal(SIGTERM, sigterm_handler)

    try:
        main()
    except KeyboardInterrupt:
        stderr.write("ok, bye\n")
        exit(1)
    except SaySomethingError as e:
        stderr.write(f"Error: {e}")
        exit(1)


if __name__ == "__main__":
    __main__()
Enter fullscreen mode Exit fullscreen mode

Now we can add the following to the pyproject.toml file:

[tool.poetry.scripts]
say_something = "say_something.__main__:__main__"
Enter fullscreen mode Exit fullscreen mode

And finally, let's try this out:

poetry run say_something -w "hello, poetry"
Enter fullscreen mode Exit fullscreen mode

This way when someone installs your package they will have the say_something command available.

Conclusion

The pattern presented in this article is one that I use all the time. Weirdly, I never did a proper template for it but now that I wrote this article I know where to come and copy it. I hope that you found it useful!

Top comments (3)

Collapse
 
luciferchase profile image
Lucifer Chase

Brilliant article. Clearly written and very helpful. Thank for writing this Remy 😁

Collapse
 
ldrscke profile image
Christian Ledermann • Edited

What is your reason to prefer argparse over something like typer or click?

Collapse
 
xowap profile image
Rémy 🤖

I never needed more and it's in the standard lib :)