Make it work on both python2 and python3 (Fixes #4)