西宝科学技术 2024-05-23

为什么写 Python 脚本时，一定要加上这个？

01 初识 if name == ‘main’

首先，我们来看看这行代码：

if __name__ == '__main__':

# 程序的主入口

main()

简简单单的一行代码，却蕴藏着 Python 的设计哲学和编程智慧。它的作用是让某些代码仅在该模块作为主程序运行时执行，而在被其他模块导入时不执行。

02 name 和 main 是什么？

__name__ 变量

每个 Python 模块都有一个内置的属性 __name__，它的值是模块的名字。如果一个模块是被导入的，__name__ 的值通常是模块的文件名，不包括路径和文件扩展名。

__main__ 值

当一个模块被直接运行时，Python 会将特殊变量 __name__ 赋值为 ‘__main__’。这意味着，如果你运行一个脚本文件，__name__ 将被设置为 ‘__main__’。而如果该模块是被导入的，那么 __name__ 将被设置为模块的名字。

03 实战讲解

示例 1：没有 if __name__ == ‘__main__’

# example.py

print("This will always be executed")

def main():

print("This will also be executed")

main()

运行 example.py：

$ python example.py

This will always be executed

This will also be executed

现在，我们创建另一个脚本来导入 example.py：

# test_import.py

import example

运行 test_import.py：

$ python test_import.py

This will always be executed

This will also be executed

不论是直接运行 example.py，还是通过 test_import.py 导入，example.py 里的所有代码都会被执行。这显然不是我们想要的结果。

示例 2：使用 if __name__ == ‘__main__’

我们来改进一下 example.py：

# example.py

print("This will always be executed")

def main():

print("This will be executed only if run directly")

if __name__ == '__main__':

main()

再次运行 example.py：

$ python example.py

This will always be executed

This will be executed only if run directly

然后运行 test_import.py：

$ python test_import.py

This will always be executed

我们发现，当 example.py 被导入时，main() 函数不会被执行。这样就达到了区分直接运行和导入执行的目的。

04 实际应用场景

测试代码

在开发过程中，常常需要测试代码段。通过 if __name__ == ‘__main__’，可以方便地在模块底部添加测试代码，而不用担心它们在模块被导入时执行：

# math_functions.py

def add(a, b):

return a + b

def subtract(a, b):

return a - b

if __name__ == '__main__':

print("Testing add function:")

print(add(1, 2)) # Should output 3

print("Testing subtract function:")

print(subtract(5, 3)) # Should output 2

提高代码可重用性

很多时候，我们写的脚本不仅仅是一次性运行的工具，而是需要被其他模块调用的函数库。

if __name__ == ‘__main__’ 可以保证模块在导入时不会执行测试代码或其他非必要代码，提高代码的可重用性。

避免不必要的执行

当我们有一段需要耗时或影响全局状态的代码时，避免在导入时执行就显得尤为重要：

# data_analysis.py

import pandas as pd

def load_data(filepath):

data = pd.read_csv(filepath)

return data

if __name__ == '__main__':

filepath = 'data.csv'

data = load_data(filepath)

print(data.head())

在导入 data_analysis.py 时，不会自动加载数据文件，只有调用模块的相应方法才会加载数据，避免了不必要的资源占用。

05 深入理解和更多用法

使用 argparse 解析命令行参数

在写脚本工具时，常常需要解析命令行参数。我们可以结合 if __name__ == ‘__main__’ 来实现这一功能：

# script.py

import argparse

def main():

parser = argparse.ArgumentParser(description="A sample command-line tool")

parser.add_argument('name', type=str, help='Your name')

args = parser.parse_args()

print(f'Hello, {args.name}!')

if __name__ == '__main__':

main()

通过命令行运行脚本：

$ python script.py Alice

Hello, Alice!

使用 unittest 进行单元测试

在模块中添加单元测试，用 if __name__ == ‘__main__’ 来运行测试：

# test_math_functions.py

import unittest

from math_functions import add, subtract

class TestMathFunctions(unittest.TestCase):

def test_add(self):

self.assertEqual(add(1, 2), 3)

def test_subtract(self):

self.assertEqual(subtract(5, 3), 2)

if __name__ == '__main__':

unittest.main()

运行测试：

$ python test_math_functions.py

使用 multiprocessing 创建子进程

当我们需要并行处理时，使用 multiprocessing 模块时也需要注意 if __name__ == ‘__main__’ 的使用：

# parallel_script.py

import multiprocessing

def worker(num):

print(f'Worker: {num}')

if __name__ == '__main__':

jobs = []

for i in range(5):

p = multiprocessing.Process(target=worker, args=(i,))

jobs.append(p)

p.start()

如果不加 if __name__ == ‘__main__’，在 Windows 系统上运行会导致无限递归创建子进程，最终导致崩溃。