Updated code for documentation