Tidy up python code re #6742

This commit was needed only because I made the previous one without saving.

Signed-off-by: Karl Palmen <karl.palmen@stfc.ac.uk>