Add support for Python 2