How to remove switch statements with polymorphism

In this post we will use polymorphism to remove switch statements. Inappropriately using switch statements is a sign of code smell.

By Subash Adhikari · 19 December 2017 · 3 min read

• Architecture

Share on XCopy link

Remove switch statements using polymorphism#

In this post we'll use polymorphism to get rid of switch statements. Using switch statements in the wrong places is a code smell. Robert Martin says in his book Clean Code that he limits himself to one switch statement per object type.

We'll use a Customer class as the example. First we'll implement its methods with switch statements. Then we'll refactor it with polymorphism to improve the structure.

public class Customer {
 
  public enum Membership {
  	Bronze, Silver, Gold
  }
 
  public Membership membership;
 
  public Customer(Membership membership) {
  	this.membership = membership;
  }
 
  public void setMembership(Membership membership) {
    this.membership = membership;
  }
 
  public Membership getMembership() {
    return membership;
  }
 
  public double getMonthlyPrice() {
    switch(membership) {
    	case Bronze:
      		return 10;
      	case Silver:
      		return 20;
      	case Gold:
      		return 30;
      	default:
      		throw new IllegalArgumentException("Invalid membership!!");
    }
  }
}

Each Customer has a membership type. The getMonthlyPrice method works out the customer's monthly membership cost based on their level. So far so good. But now we've been asked to add a way to calculate the points a customer earns. Customers earn points based on their membership level and how much they spend.

public double getPointsEarned(double amountSpent) {
  switch(membership) {
    case Bronze:
    	return amountSpent * 1;
    case Silver:
    	return amountSpent * 2;
    case Gold:
    	return amountSpent * 3;
    default:
    	throw new IllegalArgumentException("Invalid membership!!");
  }
}

Now the problem is obvious. We have duplicated switch statements. Our example is really simple, but a real application would have much more complex logic in these calculations.

The first idea might be to use inheritance and add BronzeCustomer, SilverCustomer and GoldCustomer, each extending an abstract Customer and implementing its own version of getMonthlyPrice and getPointsEarned. But inheritance won't work here, because the membership level can change at runtime after the object is created. A Customer can be upgraded or downgraded to another level at any time.

We can solve this with the State Pattern. We'll start the refactor by introducing a Member interface and delegating the price and points calculations to it. Then we'll create BronzeMember, SilverMember and GoldMember, which all implement Member.

public interface Member {
  double getMonthlyPrice();
  double getPointsEarned(double amountSpent);
}
 
class BronzeMember implements Member {
  public double getMonthlyPrice() {
    return 10;
  }
 
  public double getPointsEarned(double amountSpent) {
    return amountSpent * 1;
  }
}
 
class SilverMember implements Member {
  public double getMonthlyPrice() {
    return 20;
  }
 
  public double getPointsEarned(double amountSpent) {
    return amountSpent * 2;
  }
}
 
class GoldMember implements Member {
  public double getMonthlyPrice() {
    return 30;
  }
 
  public double getPointsEarned(double amountSpent) {
    return amountSpent * 3;
  }
}

Next we update the Customer class to use the Member interface. We add a private member property, and Customer delegates the calculations to it. When the membership is set, the member instance is updated to match the membership type. The constructor now uses setMembership too. Here's the final Customer class.

public class Customer {
 
  public enum Membership {
  	Bronze, Silver, Gold
  }
 
  private Membership membership;
  private Member member;
 
  public Customer(Membership membership) {
  	this.setMembership(membership);
  }
 
  public void setMembership(Membership membership) {
    this.membership = membership;
 
    switch(membership) {
      case Bronze:
      	this.member = new BronzeMember();
      	break;
      case Silver:
      	this.member = new SilverMember();
      	break;
      case Gold:
      	this.member = new GoldMember();
      	break;
      default:
      	throw new IllegalArgumentException("Invalid membership!!");
    }
  }
 
  public Membership getMembership() {
    return membership;
  }
 
  public double getMonthlyPrice() {
    return this.member.getMonthlyPrice();
  }
 
  public double getPointsEarned(double amountSpent) {
  	return this.member.getPointsEarned(amountSpent);
  }
}

We've used polymorphism to refactor away the switch statements in Customer. The logic for calculating membership cost and points now lives in its own class for each level. When a new membership level comes along, we just update setMembership to support it. Best of all, none of the client classes that use Customer will notice anything changed. They keep working as before.

Thanks for reading. You can also follow along on X or with RSS.