4个帮你编写完美Python脚本的技巧

声明：

1. 译自How to Write Perfect Python Command-line Interfaces --Learn by Examples
2. 本文会基于例程介绍4个完善Python脚本的技巧

四个技巧

1. 命令行参数解析中设置默认值
2. 异常处理
3. 帮助/说明文档
4. 进度条控件

从一个例程开始

Python脚本

何为Python脚本？
我们有一个.py文件，即python程序。这个程序实现了一个独立的功能，需要在命令行中运行以发挥它的作用。

例程

在密码学领域，有一个最古老、最经典的编解码算法——凯撒密码，我们编写该算法的程序，并以此为例介绍以上4个技巧。

我们的加密函数如下：

def encrypt(plaintext, key):
	"""
	:param plaintext: 源文本
	:param 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

为了对源文本加密，我们需要通过在调用以上加密函数时，需要输入参数，参数来自于命令行。因此，我们需要编写一些其他的代码来获取命令行中的参数输入，然后完成对加密函数encrypt()的调用。

Python中提供了sys.argv处理命令行参数。当我们在命令行运行Python脚本时，sys.argv是一个list，包含我们在命令行中输入的所有参数和参数值。比如，我们输入：python caesar_script.py --key 23 --decrypt my secret message，则打印sys.argv可得：

>>> print(sys.argv)
['caesar_script.py', '--key', '23', '--decrypt', 'my', 'secret', 'message']

现在，我们编写一些代码，用来解析sys.argv。

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]
			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()

上面的代码有点长，因为我们在解析sys.argv的同时，还考虑到了上文提到的4个技巧中的3点：

1. 在解析命令行参数时，别忘了设置默认值(key)。
2. 异常处理（如果有参数没有输入，或者有无效参数时，要进行报错或相应处理）
3. 要有相应的帮助文档（如果报错了，要给出建议，或者是给出该脚本正确的使用方法）

使用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中的argparse模块能自动地解析命令行参数，并且处理相应的输入异常，比如：

>>> 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

关于argparse更多的用法请参考Python的官方文档。

进度条控件

当Python脚本中有遍历操作时，使用进度条控件可以让我们知道程序运行的进度。Python中有tqdm模块提供了进度条控件，使用方式如下：

from tqdm import tqdm

for key in tqdm(range(26)):
	pass