Adding @pfreixes suggestion on __init__.py and imports