Simplify and improve code of last commit