super pythonic refactor

- Fixes #3