Minor tweak to extension