Better fix for #7 to keep original behavior