使用click创建完美的Python命令行程序

用户8949263

发布于 2022-11-07 12:56:12

1K0

发布于 2022-11-07 12:56:12

Python程序员的主要工作是写命令行程序，即直接在终端运行的脚本。随着项目规模增长，我们希望创建有效的命令行接口，通过提供不同的参数，解决不同的问题，而不是每次都修改源代码。

Click库是一个非常高效的命令行工具，能够帮助我们快速创建完美的命令行接口，小编认为这是每个Python程序员都应该掌握的工具。
作为Python开发人员，我们经常编写命令行程序。例如，在我的数据科学项目中，我会在终端运行多个脚本来训练模型并评估算法的准确性。提高生产率的一个方法是定义简单和直接的命令行程序接口，对于多人参与的项目而言更是如此。

侧重于对象的响应性，利用对象之间的交互实现更强大的功能。

为了实现这一目标，我总结了四条原则，希望对大家有所帮助：

命令行参数应提供默认值
处理所有可能的参数错误，包括缺少参数，数据类型错误，无法找到文件等
撰写完善的文档，解释参数的含义以及如何设置
使用进度条显示长时间运行的任务

一个简单的例子

让我们将这些规则应用于一个具体的案例：一个使用Caesar cipher加密和解密消息的脚本。

假设我们编写了一个encrypt函数，如下所示。现在要创建一个的脚本来加密和解密消息。

脚本允许用户选择：模式（加密或解密），密钥。前者的默认值是加密，后者的默认值是1。这一切都通过命令行参数实现。

def encrypt(plaintext, key):
    cyphertext = ''
    for character in plaintext:
        if character.isalpha():
            number = ord(character)
            number += key
            if character.isupper():
                if number > ord('Z'):
                    number -= 26
                elif number < ord('A'):
                    number += 26
            elif character.islower():
                if number > ord('z'):
                    number -= 26
                elif number < ord('a'):
                    number += 26
            character = chr(number)
        cyphertext += character

    return cyphertext

初学者方法：sys.argv

脚本需要先获取命令行参数的值，让我们先用最简单的sys.argv实现。

sys.argv是一个列表，包含了用户在运行脚本时输入的所有参数（包括脚本名字本身）。

在终端输入以下指令：

> python caesar_script.py --key 23 --decrypt my secret message
pb vhfuhw phvvdjh

sys.argv列表包括：

['caesar_script.py', '--key', '23', '--decrypt', 'my', 'secret', 'message']

为了获得参数值，需要遍历参数列表，寻找一个 '--key' （或 '-k' ）来获取密钥值，并寻找一个 '--decrypt' 获取模式。

import sys

from caesar_encryption import encrypt


def caesar():
    key = 1
    is_error = False

    for index, arg in enumerate(sys.argv):
        if arg in ['--key', '-k'] and len(sys.argv) > index + 1:
            key = int(sys.argv[index + 1])
            del sys.argv[index]
            del sys.argv[index]
            break

    for index, arg in enumerate(sys.argv):
        if arg in ['--encrypt', '-e']:
            del sys.argv[index]
            break
        if arg in ['--decrypt', '-d']:
            key = -key
            del sys.argv[index]
            break

    if len(sys.argv) == 1:
        is_error = True
    else:
        for arg in sys.argv:
            if arg.startswith('-'):
                is_error = True

    if is_error:
        print(f'Usage: python {sys.argv[0]} [ --key <key> ] [ --encrypt|decrypt ] <text>')
    else:
        print(encrypt(' '.join(sys.argv[1:]), key))

if __name__ == '__main__':
    caesar()

代码遵循我们一开始提出的原则：

有一个默认键值和一个默认模式

处理基本错误（不提供输入文本或未知参数）

在参数错误或在不带参数的情况下调用脚本时，打印简洁的提示信息

> python caesar_script_using_sys_argv.py
Usage: python caesar.py [ --key <key> ] [ --encrypt|decrypt ] <text>

但是这个版本的脚本相当长（39行，不包括加密函数），而且代码非常丑陋。

是否有更好的方法来解析命令行参数？

进入argparse

argparse是用于解析命令行参数的Python标准库模块。

修改脚本，使用argparse解析命令行参数：

import argparse

from caesar_encryption import encrypt


def caesar():
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-e', '--encrypt', action='store_true')
    group.add_argument('-d', '--decrypt', action='store_true')
    parser.add_argument('text', nargs='*')
    parser.add_argument('-k', '--key', type=int, default=1)
    args = parser.parse_args()

    text_string = ' '.join(args.text)
    key = args.key
    if args.decrypt:
        key = -key
    cyphertext = encrypt(text_string, key)
    print(cyphertext)

if __name__ == '__main__':
    caesar()

代码仍然遵守我们提出的原则，并且比手动解析命令行参数提供更精确的文档和更具交互性的错误处理。

> python caesar_script_using_argparse.py --encode My message

usage: caesar_script_using_argparse.py [-h] [-e | -d] [-k KEY] [text [text ...]]
caesar_script_using_argparse.py: error: unrecognized arguments: --encode

> python caesar_script_using_argparse.py --help

usage: caesar_script_using_argparse.py [-h] [-e | -d] [-k KEY] [text [text ...]]positional arguments:
  text

optional arguments:
  -h, --help         show this help message and exit
  -e, --encrypt
  -d, --decrypt
  -k KEY, --key KEY

脚本的第7行到第13行代码定义了命令行参数，但它们不是很优雅：太冗长且程序化，我们可以用更紧凑和声明性的方式完成。

使用click创建更好的命令行接口

幸运的是有一个三方库click用于创建命令行接口，它不仅提供比argparse更多的功能，而且代码风格更漂亮。用click替换argparse，继续优化脚本。

import click

from caesar_encryption import encrypt

@click.command()
@click.argument('text', nargs=-1)
@click.option('--decrypt/--encrypt', '-d/-e')
@click.option('--key', '-k', default=1)
def caesar(text, decrypt, key):
    text_string = ' '.join(text)
    if decrypt:
        key = -key
    cyphertext = encrypt(text_string, key)
    click.echo(cyphertext)

if __name__ == '__main__':
    caesar()

注意，命令行参数和选项都在装饰器中声明，这使得它们可以作为函数的参数直接访问。

让我们仔细分析上面的代码：

nargs定义了命令行参数接收的值的数量，默认值为1，nargs=-1允许提供任意数量的单词。

--encrypt/--decrypt定义互斥的选项，最终以布尔值传递给程序。

click.echo是click库提供的基础函数，功能类似于print，但提供更强大的功能，例如调整打印到控制台的文本的颜色。

从本地文件读取输入

命令行参数接收的值是将被加密的最高机密消息，所以如果要求用户直接在终端中输入纯文本，可能会引发安全顾虑。

一种更安全的方法是使用隐藏提示，或者从本地文件读取文本，这对于长文本来说更加实用。

这个想法同样适用于输出：用户可以将其保存到文件中，或者在终端中打印出来。让我们继续优化脚本。

import click

from caesar_encryption import encrypt

@click.command()
@click.option(
    '--input_file',
    type=click.File('r'),
    help='File in which there is the text you want to encrypt/decrypt.'
         'If not provided, a prompt will allow you to type the input text.',
)
@click.option(
    '--output_file',
    type=click.File('w'),
    help='File in which the encrypted / decrypted text will be written.'
         'If not provided, the output text will just be printed.',
)
@click.option(
    '--decrypt/--encrypt',
    '-d/-e',
    help='Whether you want to encrypt the input text or decrypt it.'
)
@click.option(
    '--key',
    '-k',
    default=1,
    help='The numeric key to use for the caesar encryption / decryption.'
)
def caesar(input_file, output_file, decrypt, key):
    if input_file:
        text = input_file.read()
    else:
        text = click.prompt('Enter a text', hide_input=not decrypt)
    if decrypt:
        key = -key
    cyphertext = encrypt(text, key)
    if output_file:
        output_file.write(cyphertext)
    else:
        click.echo(cyphertext)

if __name__ == '__main__':
    caesar()

由于脚本变得更复杂，我们创建了参数文档（通过定义click.option装饰器的help参数实现），详细解释参数的功能，效果如下。

> python caesar_script_v2.py --help

Usage: caesar_script_v2.py [OPTIONS]

Options:
  --input_file FILENAME          File in which there is the text you want to encrypt/decrypt. If not provided, a prompt will allow you to type the input text.
  --output_file FILENAME         File in which the encrypted/decrypted text will be written. If not provided, the output text will just be printed.
  -d, --decrypt / -e, --encrypt  Whether you want to encrypt the input text or decrypt it.
  -k, --key INTEGER              The numeric key to use for the caesar encryption / decryption.
  --help                         Show this message and exit.

我们有两个新的参数input_file和output_file，类型是click.File，click会用正确的模式打开文件并处理可能发生的错误。例如找不到文件：

> python caesar_script_v2.py --decrypt --input_file wrong_file.txt

Usage: caesar_script_v2.py [OPTIONS]

Error: Invalid value for "--input_file": Could not open file: wrong_file.txt: No such file or directory

如果未提供input_file，则我们用click.prompt，在命令行创建提示窗口，让用户直接输入文本，该提示对于加密模式将是隐藏的。效果如下：

> python caesar_script_v2.py --encrypt --key 2

Enter a text: **************

yyy.ukectc.eqo

假设你是一名黑客：想要解密一个用凯撒加密过的密文，但你不知道秘钥是什么。最简单的策略就是用所有可能的秘钥调用解密函数 25 次，阅读解密结果，看看哪个是合理的。但你很聪明，而且也很懒，所以你想让整个过程自动化。确定解密后的 25 个文本哪个最可能是原始文本的方法之一，就是统计所有这些文本中的英文单词的个数。这可以使用 PyEnchant 模块实现：

import click
import enchant

from caesar_encryption import encrypt

@click.command()
@click.option(
    '--input_file',
    type=click.File('r'),
    required=True,
)
@click.option(
    '--output_file',
    type=click.File('w'),
    required=True,
)
def caesar_breaker(input_file, output_file):
    cyphertext = input_file.read()
    english_dictionnary = enchant.Dict("en_US")
    max_number_of_english_words = 0
    for key in range(26):
        plaintext = encrypt(cyphertext, -key)
        number_of_english_words = 0
        for word in plaintext.split(' '):
            if word and english_dictionnary.check(word):
                number_of_english_words += 1
        if number_of_english_words > max_number_of_english_words:
            max_number_of_english_words = number_of_english_words
            best_plaintext = plaintext
            best_key = key
    click.echo(f'The most likely encryption key is {best_key}. It gives the following plaintext:\n\n{best_plaintext[:1000]}...')
    output_file.write(best_plaintext)

if __name__ == '__main__':
    caesar_breaker()

使用进度条

示例中的文本包含10^4个单词，因此该脚本需要大约5秒才能解密。这很正常，因为它需要检查所有25个秘钥，每个秘钥都要检查10^4个单词是否出现在英文字典中。

假设你要解密的文本包括10^5个单词，那么就要花费50秒才能输出结果，用户可能会非常着急。因此我建议这种任务一定要显示进度条。特别是，显示进度条还非常容易实现。下面是个显示进度条的例子：

import click
import enchant

from tqdm import tqdm

from caesar_encryption import encrypt

@click.command()
@click.option(
    '--input_file',
    type=click.File('r'),
    required=True,
)
@click.option(
    '--output_file',
    type=click.File('w'),
    required=True,
)
def caesar_breaker(input_file, output_file):
    cyphertext = input_file.read()
    english_dictionnary = enchant.Dict("en_US")
    best_number_of_english_words = 0
    for key in tqdm(range(26)):
        plaintext = encrypt(cyphertext, -key)
        number_of_english_words = 0
        for word in plaintext.split(' '):
            if word and english_dictionnary.check(word):
                number_of_english_words += 1
        if number_of_english_words > best_number_of_english_words:
            best_number_of_english_words = number_of_english_words
            best_plaintext = plaintext
            best_key = key
    click.echo(f'The most likely encryption key is {best_key}. It gives the following plaintext:\n\n{best_plaintext[:1000]}...')
    output_file.write(best_plaintext)

if __name__ == '__main__':
    caesar_breaker()

这里使用了tqdm库，tqdm.tqdm类可以将任何可迭代对象转化为一个进度条。click也提供了类似的接口来创建进度条（click.progress_bar），但我觉得它不如tqdm好用。