Refactor python logic

[flindeberg]

Python logic is more or less a hack (which ostensibly works), should be fixed.