Rework python setup code