Refactoring - Replace Type Code with Polymorphism

This type of refactoring applies when a class has a property that represent a type of some kind and the code is depending on it. This usually manifests through if/else or switch statements.

Let’s continue the example of the Karateka but simplify it a bit.


Assume the belts are of three kinds: white, blue and black. They correspond to a beginner, intermediate and, respectively, advanced level.

First we define an enumeration with the three belts:

enum Belt {
    case White, Blue, Black
}

Then, we’ll define the Karateka class that has two methods for displaying the title and the level:


class Karateka {
    var belt = Belt.White
   
    func getTitle() -> String {
        switch belt {
        case .White :
            return "White belt karateka"
        case .Blue :
            return "Blue belt karateka"
        case .Black :
            return "Black belt karateka"
        }
    }
   
    func getLevel() -> String {
        switch belt {
        case .White :
            return "Beginner karateka"
        case .Blue :
            return "Intermediate karateka"
        case .Black :
            return "Advanced karateka"          
        }
    }
}

Here is some code to test the class:

let karateka1 = Karateka()
karateka1.belt = .White

let karateka2 = Karateka()
karateka2.belt = .Blue

let karateka3 = Karateka()
karateka3.belt = .Black

karateka1.getTitle()
karateka2.getTitle()
karateka3.getTitle()

karateka1.getLevel()
karateka2.getLevel()
karateka3.getLevel()

You observe the switch statements that need to be replicated for each piece of functionality that depends on the belt color. Also, if one more belt is added, all the switch statements need to be changed.

Instead of this, we’ll create a base class and for each belt and we’ll overwrite the methods that depend on the belt type property.

class Karateka {
    var belt: Belt
   
    init() {
        belt = .White
    }
   
    func getTitle() -> String {
        return "unknown"
    }
   
    func getLevel() -> String {
        return "unknown"
    }
}

Now let’s implement one class for each belt:

class KaratekaWhiteBelt: Karateka {
    override init() {
        super.init()
        super.belt = .White
    }
   
    override func getTitle() -> String {
        return "White belt karateka"
    }
   
    override func getLevel() -> String {
        return "Beginner karateka"
    }
}

class KaratekaBlueBelt: Karateka {
    override init() {
        super.init()
        super.belt = .Blue
    }
   
    override func getTitle() -> String {
        return "Blue belt karateka"
    }
   
    override func getLevel() -> String {
        return "Intermediate karateka"
    }
}

class KaratekaBlackBelt: Karateka {
    override init() {
        super.init()
        super.belt = .Black
    }
   
    override func getTitle() -> String {
        return "Black belt karateka"
    }
   
    override func getLevel() -> String {
        return "Advanced karateka"
    }
}

And here is some code to test our new refactored functionality:

let karateka1 = KaratekaWhiteBelt()
let karateka2 = KaratekaBlueBelt()
let karateka3 = KaratekaBlackBelt()

karateka1.getTitle()
karateka2.getTitle()
karateka3.getTitle()

karateka1.getLevel()
karateka2.getLevel()
karateka3.getLevel()

Code looks more elegant now.

If we need to add more belts, all we have to do is add more subclasses and, of course, update the Belt enumeration.

All the other code stays the same.