Moved code out of __init__.py, minor pep8 fix