Fix import
Merge pull request #2 from omershelef/add-python3-support

Fix import to support python 3