Martin Nyaga

Turning edge-case if statements into polymorphism

December 1st, 2016


As your codebase gets larger and larger, if statements can be bad for your health. This is especially true if you use them in a class with a lot of logic, and you use them to handle tiny edge cases. You find yourself scattering the if statement all over the class and making the code very messy.

Let’s dive straight into some ruby code to illustrate this mess.

Imagine I have a codebase with a ParkingTicket class which determines how much a parking ticket will be, given some options.

The fee is calculated from a base fee, plus a surcharge depending on the day of the week. Imagine there’s also a lot of other logic in this class that I don’t want to focus on for this purpose.

Here is a half decent implementation of it:

class ParkingTicket
  BASE FEE = 1000

  # This receives a long list of options for
  # some other logic that I don't care about
  def initialize(opts)
    @day = opts[:day]
    @vehicle_type = opts[:vehicle_type]

    # EET stands for some kind of environmental
    # friendliness test
    @eet_passed = opts[:eet_passed]

    # more code that's not important for this purpose... 
  end

  def total_cost
    BASE_FEE + surcharge
  end

  def surcharge
    case @day
    when "Monday", "Tuesday", "Wednesday", "Friday"
      100
    when "Thursday", "Saturday", "Sunday"
      150
    end
  end
end

Now, imagine we are then told that trucks have a higher surcharge by a factor 1.5. So we do the dangerously easy thing to do and add an if statement in the surcharge method:

def surcharge
  charge =
    case @day
    when "Monday", "Tuesday", "Wednesday", "Friday"
      100
    when "Thursday", "Saturday", "Sunday"
      150
    end

  if @vehicle_type = "Truck"
    charge *= 1.5
  end

  charge
end

Neat.

Until you’re told that matatus have a higher surcharge on weekdays, by a factor of 1.7. Uncomfy, but still not too hard:

def surcharge
  charge =
    case @day
    when "Monday", "Tuesday", "Wednesday", "Friday"
      100
    when "Thursday", "Saturday", "Sunday"
      150
    end

  if @vehicle_type = "Truck"
    charge *= 1.5
  end

  if @vehicle_type = "Matatu" && weekdays.include?(@day)
    charge *= 1.7
  end

  charge
end

private
def weekdays
  ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]
end

Then you’re told that public buses that pass the environmental friendliness test (EET), only get charged the corresponding surcharge and not the base fee. Also, matatus and trucks have different base fees now of 1100 and 1200 respectively.

Okay, yuck, but here goes:

# get rid of the base fee constant
def base_fee
  if @vehicle_type == "Matatu"
    1100
  elsif @vehicle_type == "Truck"
    1200
  else
    1000
  end
end

def total_cost
  if @vehicle_type == "Public Bus" && @eet_passed
    surcharge
  else
    base_fee + surcharge
  end
end

You can see where this is going. This class is going to get very messy. The ParkingTicket class is not only keeping track of the core logic and behaviour of a ticket, but also keeping track of all the weird rule exceptions for specific vehicle types.

And if you were given a completely new set rules for just matatus, you’d start by digging through the code hunting for if statements, and end up slitting your wrists. This is poor object oriented design.

A good way to fix this is to use polymorphism.

What is polymorphism?

In very simple terms, polymorphism means allowing a certain thing of interest to take many forms.

For example, when you write a + b in ruby, it has no restrictions, per se, about what thing of type a is. It only expects that a responds to :+ and that the implementation of :+ in a can accept b as a parameter.

So a can be an integer, a float, a string, or an array, it doesn't matter. Thing a can take many different forms i.e. is polymorphic. The only thing we care about is that it responds to the method we're trying to call on it.

Here’s another example:

class Animal
  def initialize(name)
    @name = name
  end

  def make_sound
    puts "No sound"
  end
end

class Dog < Animal
  def make_sound
    puts "Woof!"
  end
end

So if elsewhere in my codebase I have some code that expects a thing that makes sound, it can be passed an animal, or a dog. It doesn’t matter.

Lets fix it!

Let’s fix the ParkingTicket class now.

We can create a number of small classes that inherit from the main ParkingTicket , each knowing their own specific rules. These classes will implement the same methods as a ParkingTicket and so can be used in place of it, wherever it is used.

First we fix up the base ParkingTicket class as follows (It’s the same as what we started with, except base_fee is a method and not a constant):

class ParkingTicket
  def initialize(opts)
    @day = opts[:day]
    @vehicle_type = opts[:vehicle_type]
    @eet_passed = opts[:eet_passed]
    ... 
  end

  def base_fee
    1000
  end

  def total_cost
    base_fee + surcharge
  end

  def surcharge
    case @day
    when "Monday", "Tuesday", "Wednesday", "Friday"
      100
    when "Thursday", "Saturday", "Sunday"
      150
    end
  end
  ... 
end

Then I’ll start with a Truck. We have two rules for a truck. We will create a class for them and override the methods we need to, to implement the rules.

  1. Trucks have a higher surcharge that’s higher by a factor 1.5
  2. Truck has a higher base_fee of 1200
class TruckParkingTicket < ParkingTicket
  def base_fee
    1200
  end

  def surcharge
    super * 1.5
  end
end

Nice and simple. Let’s move on to a Matatu. We have two rules for a matatu.

  1. Matatus have a higher surcharge on weekdays
  2. Matatus have a higher base_fee of 1100
class MatatuParkingTicket < ParkingTicket
  def base_fee
    1100
  end

  def surcharge
    weekdays.inlcude?(@day) ? super * 1.5  : super
  end

  private
    def weekdays
      ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday"]
    end
end

Lastly, we have a public bus, which only has one rule

  1. public buses that pass the environmental friendliness test (EET), only get charged the surcharge and not the base fee
class PublicBusParkingTicket < ParkingTicket
  def total_cost
    @eet_passed == true ? surcharge : super
  end
end

So far so good. We have managed to push the rules to little objects that clearly express the differences and can be used in context. But how do we use them?

Meet the TicketFactory

We need a way to choose what parking ticket to use. This is a good problem, because we have taken the littered if statements and turned them into just one point of decision.

We can use the factory pattern to solve this. A factory class is a class whose job is to instantiate objects of various classes using some rules (hence the name " FactoryGirl " for the gem). Here is an example of a factory class for the tickets.

class TicketFactory
  def self.ticket_for(opts)
    class_for(opts[:vehicle_type]).new(opts)
  end

  def self.class_for(vehicle_type)
    case vehicle_type
    when 'Truck'
      TruckParkingTicket
    when 'Matatu'
      MatatuParkingTicket
    when 'Public Bus'
      PublicBusParkingTicket
    else
      ParkingTicket
    end
  end
end

Then whenever we need to create a ticket, we call it like so:

parking_ticket = TicketFactory.ticket_for(opts)
# where opts is the hash of options

We can further refactor the TicketFactory as follows, to remove the pesky case statement:

class TicketFactory
  TICKET_CLASSES = {
    'Truck' => TruckParkingTicket,
    'Matatu' => MatatuParkingTicket,
    'Public Bus' => PublicBusParkingTicket
  }

  def self.ticket_for opts
    # .fetch takes a second argument as a default if the 
    # key is not found
    TICKET_CLASSES.fetch(
      opts[:vehicle_type],
      ParkingTicket
    ).new(opts)
  end
end

Now, if matatu rules change, it’s easy to implement. Just go to the MatatuParkingTicket class and make the changes there. To add new specific rules for Private Cars for example, we just add a PrivateCarParkingTicket and put the rules there, then update our ticket factory's TICKET_CLASSES with the new entry. No more hair pulling shotgun surgery!

Extra Rails Goodness

If we really want to over-engineer this in a rails environment, we can take a “Convention over configuration” approach.

We can see the general trend of the vehicle type data here is fairly consistent. We don’t expect that some day the vehicle type will be a number.

If we can make this assumption, then we can condense the TicketFactory even further. We can generate the class to be instantiated using the actual string given by the vehicle_type.

class TicketFactory
  def self.ticket_for(opts)
    class_for(opts[:vehicle_type]).new(opts)
  end

  # Method to auto generate class from a vehicle_type
  # e.g "Private car" => PrivateCarParkingTicket
  #     "Matatu" => MatatuParkingTicket
  def self.class_for(vehicle_type)
    prefix = vehicle_type.underscore.parameterize("_").camelize
    "#{prefix}ParkingTicket".constantize

    # If the class was not found
    rescue NameError
      ParkingTicket
    end
  end
end

This will turn “Matatu” into MatatuParkingTicket and "Public bus" into PublicBusParkingTicket .

The advantage of this is that now we don’t even need to touch the TicketFactory to add new rules. Just add a new class, named appropriately, and add your rules there. No need to touch any existing code!

tl;dr

A bunch of if statements to handle edge cases is probably bundling responsibility and hiding some object(s) that should exist. Look for those objects, extract them into polymorphic classes and use the factory pattern to instantiate them.