Python 模拟：如果设置了 sys.argv 参数，则用另一个函数替换函数

Question

我有一段代码看起来像这样：

if args.module == 'omega':
    with mock.patch.object(quantum_entangler, 'sort_by_dirac', sort_by_heisenberg):
        quantum_entangler.main(atom_no)
else:
    quantum_entangler.main(atom_no)

这基本上让我在命令行调用中使用

--module

选项设置排序功能。

我想知道是否有一种方法可以只使用没有 if/else 块的

mock.patch.object()

行来实现它。

Answer 1

您可以使用字典来排序功能

想法是创建一个字典，其中包含对排序函数的引用。

以下代码可在名为

test.py

的文件中找到。要执行

test.py

我必须通过命令行传递2个参数：

• --module

  具有值之一

  omega

  |

  dirac

  
  NOTE. 我已经定义了值

  dirac

  以通过函数 sort_by_dirac() 执行正常排序，而通过

  omega

  选择通过函数 sort_by_heisenberg() 的排序，就像您在代码中所做的那样

• --atom_no

  可以是任何值（它用于传递一些东西给模块的main()函数

  quantum_entangler.py

  ）：

test.py

的代码：

from unittest import mock
import argparse
import quantum_entangler
from quantum_entangler import sort_by_heisenberg, sort_by_dirac

parser = argparse.ArgumentParser()
parser.add_argument("--module")
parser.add_argument("--atom_no")
args = parser.parse_args()

atom_no = args.atom_no

# THIS IS THE DICTIONARY of references
sortFunction = {'omega': sort_by_heisenberg, 'delta': sort_by_dirac}

# ---> here you find patch.object() function WITHOUT if/else as you request
# in the question
with mock.patch.object(quantum_entangler,'sort_by_dirac',sortFunction[args.module]):
    quantum_entangler.main(atom_no)

执行test.py

为了测试代码，我编写了一个具有以下 3 个功能的模块

quantum_entangler.py

：

def sort_by_dirac(atom_no):
    print(f"execute sort_by_dirac({atom_no})")
    
def sort_by_heisenberg(atom_no):
    print(f"execute sort_by_heisenberg({atom_no})")
    
def main(atom_no):
    sort_by_dirac(atom_no)

用

test.py

执行

--module dirac

：

python test.py --atom_no 10 --module dirac

# output
execute sort_by_dirac(10)

用

test.py

执行

--module omega

：

python test.py --atom_no 10 --module omega

# output
execute sort_by_heisenberg(10)

Answer 2

你可以尝试这样的事情：

sort_function = sort_by_heisenberg if args.module == 'omega' else sort_by_dirac

with mock.patch.object(quantum_entangler, 'sort_by_dirac', sort_function):
    quantum_entangler.main(atom_no)