Refactoring to split main.py into pieces