1

I have a snippet of code that looks like this:

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)

Which basically lets me set sort function with --module option in the command line call.

I was wondering if there's a way to implement this with just the mock.patch.object() line without the if/else block.

frankfalse
  • 1,553
  • 1
  • 4
  • 17
Oneiros
  • 143
  • 1
  • 13

3 Answers3

1

You can use a dictionary of references to the sorting function

The idea is to create a dictionary which contains references to the sorting functions.

The following code is found in a file called test.py. To execute test.py I have to pass it 2 arguments by command line:

  • --module with one of the value omega | dirac
    NOTE. I have defined the value dirac to execute the normal sorting by the function sort_by_dirac(), while by omega is selected the sorting by the function sort_by_heisenberg() as you already do in your code
  • --atom_no which can be of any value (it is used to pass something to main() function of the module quantum_entangler.py):

Code of 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)

Execution of test.py

To test the code I have written a module quantum_entangler.py with the following 3 functions:

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)

Execution of test.py with --module dirac:

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

# output
execute sort_by_dirac(10)

Execution of test.py with --module omega:

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

# output
execute sort_by_heisenberg(10)
frankfalse
  • 1,553
  • 1
  • 4
  • 17
1

Fully aware that this is out-of-scope for the current question, but:

If you have the capacity to influence the API of the quantum_entangler.main function, it might make it both easier to reason-about and easier to test if you make that function a parameter of main instead of trying to patch it in.

Borrowing some of the scaffolding code from the poster above, this could look like:

import argparse

def sort_by_heisenberg(atom_no):
    ...

def sort_by_dirac(atom_no):
    ...

def main(atom_no, sort_fn = sort_by_dirac):
    print(f"executing {sort_fn.__name__}({atom_no})")
    sort_fn(atom_no)


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

sort_fn = sort_by_heisenberg if args.module == "omega" else sort_by_dirac

main(args.atom_no, sort_fn)

With this as the result:

/ # python file.py --atom-no 10
executing sort_by_dirac(10)
/ # python file.py --atom-no 10 --module omega
executing sort_by_heisenberg(10)

A benefit to that is that it wouldn't be a breaking change for anyone currently consuming the API, and is easier to debug since otherwise it wouldn't be clear that the sorting function could be different under different contexts

vellista
  • 96
  • 2
0

you can try something like:

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)
yogev
  • 189
  • 5