如何无需if/else块，用mock.patch.object实现命令行驱动的函数替换？

当然可以去掉if/else分支，直接用mock.patch.object的参数逻辑搞定。核心就是给patch.object的new参数传一个条件判断结果——需要替换时传目标函数，不需要替换时传原对象（或者None，不过传原对象更严谨）。

修改后的代码如下：

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

逻辑说明： #

• 当args.module == 'omega'时，new参数是sort_by_heisenberg，这时候patch会把quantum_entangler.sort_by_dirac替换成这个函数，和原if分支的行为一致。
• 当条件不满足时，new参数是quantum_entangler.sort_by_dirac自己，相当于patch没有做任何替换，上下文里调用quantum_entangler.main()的效果和原else分支完全一样。

如果图省事，也可以把else分支的参数换成None，这时候mock会默认创建一个MagicMock，但如果你的代码不需要这个mock对象，传原对象更贴合原逻辑。

内容的提问来源于stack exchange，提问作者Oneiros