Cleaning up main.py