Re-added this code accidentally removed in the last commit