Fix SyntaxError: 'return' with argument inside generator