Fix missing import for python 2
Merge pull request #2 from vaz-ar/master

Fix missing import for python 2