Proper monkey patching
+ one if/else simplified