Update code.rb

Refactored code