Cleaned up code in setup.py