Made some cleanup tweaks, thanks to Pycharm