ENH: Fixed Steve's code review comments