Improved python code