12
\$\begingroup\$

I'm building a command-line tool that lets you use Python's secrets module. Here are some commands you can run with this tool: sct -tokenhex 8, sct -tokenurlsafe 16, sct -tokenbytes 16

I added a few commands but want to know what a better way to refactor this would be:

import secrets
import sys


if __name__ == '__main__':
    if len(sys.argv) == 0:
        print('Usage: sct [OPTION]...')
        exit(1)
    
    args = sys.argv[1:]
    
    for i in range(0, len(args) - 1):
        arg = args[i]

        if arg[0] == '-':
            if arg[1:] == 'tokenbytes':
                print(secrets.token_bytes(int(args[i+1])))
                exit(0)
            elif arg[1:] == 'tokenhex':
                print(secrets.token_hex(int(args[i+1])))
                exit(0)
            elif arg[1:] == 'tokenurlsafe':
                print(secrets.token_urlsafe(int(args[i+1])))
                exit(0)
        else:
            print('Invalid syntax')
            exit(2)


    print(args)
\$\endgroup\$
2
  • 2
    \$\begingroup\$ Do you really want to print raw bytes? Is this to be consumed in an interactive shell or from another script - if the former, I would expect something like a hex representation instead. \$\endgroup\$
    – Reinderien
    Commented May 29 at 0:35
  • \$\begingroup\$ Neat. Is this just for curiosity or do you want to contribute your code to the module? You know that command-line python-secrets already exists? maybe check out its command-line feature-list. \$\endgroup\$
    – smci
    Commented May 30 at 2:06

4 Answers 4

16
\$\begingroup\$

UX for the CLI

some commands you can run with this tool: sct -tokenhex 8, ...

That's nice. But wouldn't you like to also expose def tokenhex() as a public library routine that app code can call into?

The single - dash is nonstandard and confusing. As written it looks like sct -t -o -k -e -n -h -e -x, where presumably a double -e -e is harmless. Given that this is a multi-character option, please prefix it with a double -- dash.

guarded code

if __name__ == '__main__':

Oh, dear! I fear you may have misunderstood why we conventionally use that guard.

Typically a module will def a bunch of functions, or perhaps a class, and then (brief!) code after the guard will invoke a function or two, perhaps as a CLI or as a demo. And due to the guard, a test suite or other caller can safely import the module without any unfortunate side effects.

But here, everything is within the guard, so there's nothing to import. Please define three functions. They will comprise your Public API, and will be super thin wrappers which merely do an int() conversion. Please define a fourth function, perhaps def main():, containing your CLI logic. I assume that as the weeks go by you will likely write additional code which will be a bigger value add than the current thin wrapper.

cracking argv

You're doing it the hard way,
in a manner that is user unfriendly, with no --help diagnostic message.

The argparse module has been around for a long time to address these needs. But for your use case, you'd be much better off with import typer. Then some trivial type annotations on your def main( ... ): function would give you parsing + help "for free".

    if len(sys.argv) == 0:

Are you sure that's what you intended? Maybe you meant to apply that test to args instead? Recall that argv[0] will be the name of your script, sct. No need to crack argv yourself -- let typer do the heavy lifting. It will even exit with appropriate error code in the event that user failed to supply appropriate args.

    for i in range(0, len(args) - 1):

nit: You can omit that zero start, as it is the default. No biggie.

Calling exit(0) is seldom the right thing to do. Usually it indicates we want the code packaged up within a def function, and should instead be doing a return from that function.

\$\endgroup\$
7
\$\begingroup\$
print(secrets.token_bytes(int(args[i+1])))

doesn't particularly work, or at least not in a form I'd expect. It's possible for a script to want to call this and read from stdout in binary format, but that's not what your code does - instead, it basically prints the repr of some bytes, so some will be encoded and others won't. I imagine that a caller would either want raw binary, or use the hex version instead. With raw binary, you'll also want to make a decision (more carefully than I have) about how to separate different tokens - \n risks being ambiguous.

Use the built-in argparse. A caveat: unless you get much more creative, this is going to reorder command-line switches and group by type.

import argparse
import secrets
import sys
import typing


def positive(arg: str) -> int:
    value = int(arg)
    if value < 1:
        raise ValueError(f'Must pass a positive integer, not {value}')
    return value


def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Generate token data from the secrets module',
    )
    parser.add_argument(
        '-b', '--bytes', type=positive, action='append', default=[],
    )
    parser.add_argument(
        '-x', '--hex', type=positive, action='append', default=[],
    )
    parser.add_argument(
        '-u', '--urlsafe', type=positive, action='append', default=[],
    )
    return parser.parse_args()


def render(
    bytes_: list[int],
    hex_: list[int],
    urlsafe: list[int],
    out: typing.TextIO,
) -> None:
    if len(bytes_) > 0:
        out.buffer.write(b'\n'.join(
            secrets.token_bytes(size)
            for size in bytes_
        ))
        out.buffer.flush()
        print(file=out)

    if len(hex_) > 0:
        print('\n'.join(
            secrets.token_hex(size) for size in hex_
        ), file=out)

    if len(urlsafe) > 0:
        print('\n'.join(
            secrets.token_urlsafe(size) for size in urlsafe
        ), file=out)


def main() -> None:
    args = get_args()
    render(bytes_=args.bytes, hex_=args.hex, urlsafe=args.urlsafe, out=sys.stdout)


if __name__ == '__main__':
    main()
$ python 292299.py -b 1 -b 2 -u 1 -u 2 -x 3 -x 4
▬
�♠
a606bc
0f7f54b6
pQ
Yk0
\$\endgroup\$
5
\$\begingroup\$

Unreachable code

I think the following condition can never be true:

if len(sys.argv) == 0:

argv will always have a length greater than 0 because the 1st element of the list is the script name (sct).

Perhaps you meant:

if len(sys.argv) == 1:

Running sct without any agruments will then show the usage.

If that is the case, and you make that change, then the following line becomes unreachable and can be deleted:

print(args)

Usage

The usage should show some examples like you have in your question:

sct -tokenhex 8
sct -tokenurlsafe 16

DRY

The following code is repeated several times:

int(args[i+1])

It might be worth it to use a variable instead, and give the variable a more meaningful name than args, especially since there are many "arg" variables in your code:

number = int(args[i+1])

Naming

It would also be nice to give arg[1:] a more meaningful name like token.

Simpler

This line:

for i in range(0, len(args) - 1):

is more simply written without the 0:

for i in range(len(args) - 1):

Documentation

You could add usage information to a header comment docstring. Also mention what the code does in general, as well as the expected output.


Here is the code with many of the suggestions above:

if __name__ == '__main__':
    if len(sys.argv) == 1:
        print('Usage: sct [OPTION]...')
        exit(1)
    
    args = sys.argv[1:]
    
    for i in range(len(args) - 1):
        arg = args[i]

        if arg[0] == '-':
            number = int(args[i+1])
            token  = arg[1:]
            if token == 'tokenbytes':
                print(secrets.token_bytes(number))
                exit(0)
            elif token == 'tokenhex':
                print(secrets.token_hex(number))
                exit(0)
            elif token == 'tokenurlsafe':
                print(secrets.token_urlsafe(number))
                exit(0)
        else:
            print('Invalid syntax')
            exit(2)
\$\endgroup\$
2
  • 2
    \$\begingroup\$ This doesn't represent good use of a main guard. The whole point of a main guard is for an importer to be able to selectively call subroutines, and that isn't possible here. \$\endgroup\$
    – Reinderien
    Commented May 28 at 20:31
  • \$\begingroup\$ Why loop through the args list with range, when you can use the more Pythonic for arg in args? If you necessarily need the index, use for idx, arg in enumerate(args) \$\endgroup\$
    – TomG
    Commented Jun 1 at 9:28
2
\$\begingroup\$

Consider using docopt to create beautiful command-line interfaces

Whenever you are creating a command-line utility, consider using the docopt module. All you need to do is create a usage string and docopt will parse the usage string and interpret your arguments with very little effort or code on your part.

See:
https://pypi.org/project/docopt/
and
http://docopt.org

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.