Practical Guide to Protocol Extensions in Swift

Protocol Extensions allow you to add default implementations and computed properties to protocols. This mean, when a type conforms to a protocol, it automatically gains the functionalities provided by the protocol extension. In this article, you will learn about:

• Providing default implementations
• Conditional Extensions
• Adding New Methods
• Swift Standard Library Extensions
• Where to go next?

Let's explore some practical use cases of protocol extensions with examples.

Providing default implementations

Let say, you're building a messaging app, and you want to define a protocol for message display. You also want to provide a default implementation for a common feature, such as showing a timestamp for each message.

// Protocol for message display
protocol MessageDisplay {
    var message: String { get }
    var sender: String { get }
    var timestamp: Date { get }
    
    func displayMessage()
}

Conform MessageDisplay protocol like below:

struct TextMessage: MessageDisplay {
    var message: String
    var sender: String
    var timestamp: Date
    
    func displayMessage() {
        print("\(sender) said: \(message) and sent at: \(timestamp)")
    }
}

class ImageMessage: MessageDisplay {
    var message: String
    var sender: String
    var timestamp: Date
    var imageURL: URL
    
    init(message: String, sender: String, timestamp: Date, imageURL: URL) {
        self.message = message
        self.sender = sender
        self.timestamp = timestamp
        self.imageURL = imageURL
    }
    
    func displayMessage() {
        print("\(sender) said: \(message) and sent at: \(timestamp)")
    }
}

Each conforming type (TextMessage and ImageMessage) has to provide its own implementation of the displayMessage method. This might lead to code duplication and potential inconsistencies in how messages are displayed across different types.

Using a protocol extension to provide a default implementation for the displayMessage method, you eliminate the need for each type to reimplement the same functionality. For example:

extension MessageDisplay {
    func displayMessage() {
        print("\(sender) said: \(message) and sent at: \(timestamp)")
    }
}

This makes it easier to manage and maintain the codebase, especially when multiple types conform to the protocol and share common behavior.

Why you should implement a default bahviour with protocol extension?

• All conforming types automatically get the default behavior without having to implement it individually.
• This allows existing conforming types to adopt the new requirement without requiring immediate changes. This is particularly useful for maintaining compatibility in evolving codebases.
• It support late binding, which means that the actual implementation is determined at runtime based on the type of the instance.

Conditional Extensions

Most of the iOS applications display an alert in case of an error. In order to display an alert, you need an instance of UIViewControllerclass. Let's create an example where we use protocol extensions with conditional extensions to display error alerts using UIAlertController only for instances of UIViewController.

First, define a protocol named ErrorAlertDisplayable that outlines the behavior to display error alerts:

protocol ErrorAlertDisplayable {
    func displayErrorAlert(message: String)
}

extension ErrorAlertDisplayable where Self: UIViewController {
    
    func displayErrorAlert(message: String) {
        let alertController = UIAlertController(title: "Error Alert!", message: message, preferredStyle: .alert)
        let okAction = UIAlertAction(title: "Okay", style: .default, handler: nil)
        alertController.addAction(okAction)
        present(alertController, animated: true, completion: nil)
    }
}

In the extension, we provide the default implementation of displayErrorAlert(message:) method, but restrict it to instances of UIViewController using the where clause.

Now, create a view controller that conforms to ErrorAlertDisplayable and use the displayErrorAlert(message:) method to display error alerts:

class ViewController: UIViewController, ErrorAlertDisplayable {
    override func viewDidLoad() {
        super.viewDidLoad()
        displayErrorAlert(message: "Something went wrong.")
    }
}

Here, we've defined a protocol ErrorAlertDisplayable that provides a method to display error alerts. Then, we've used a conditional extension to provide a default implementation of the method, but only for instances of UIViewController.

This approach ensures that only instances of UIViewController can display error alerts using the UIAlertController, making the code more focused and easy to manage. Other types that do not conform to UIViewController won't be able to access this method, preventing misuse.

How conditional extensions are useful?

• It allow you to provide specialized behavior for specific types that conform to a protocol. This enables you to add tailored functionality to individual types without affecting other types that conform to the same protocol.
• You can reuse the same protocol and add distinct behaviors to different types. This reduces code duplication and promotes the DRY (Don't Repeat Yourself) principle.
• Instead of cluttering your main types with conditional logic, you can separate specialized behavior into extensions. This keeps your main types focused on their core responsibilities and makes your codebase cleaner and more understandable.

Adding New Methods

When your codebase evolves over time, you can introduce new methods or properties in protocol extensions without affecting existing conforming types.

Let say, you are building an application that simulates a task management system. You want to compare the priority of two task queues based on their number of pending tasks. You decide to use the Queue protocol to define the behavior of your task queues.

See the below example of how you can implement and use this:

protocol Queue {
    associatedtype ItemType
    var count: Int { get }
    func push(_ element: ItemType)
    func pop() -> ItemType
}

extension Queue {
    func compare<Q>(queue: Q) -> ComparisonResult where Q: Queue  {
        if count < queue.count { return .orderedAscending }
        if count > queue.count { return .orderedDescending }
        return .orderedSame
    }
}

// Define a struct for tasks
struct Task {
    let title: String
    let priority: Int
}

// Implement a TaskQueue conforming to the Queue protocol
class TaskQueue: Queue {
    typealias ItemType = Task
    private var tasks: [Task] = []
    
    var count: Int {
        return tasks.count
    }
    
    func push(_ task: Task) {
        tasks.append(task)
    }
    
    func pop() -> Task {
        return tasks.removeFirst()
    }
}

Here is the usage of above example:

// Create two task queues
var queue1 = TaskQueue()
queue1.push(Task(title: "Fix bug", priority: 3))
queue1.push(Task(title: "Implement feature", priority: 1))

var queue2 = TaskQueue()
queue2.push(Task(title: "Write documentation", priority: 2))
queue2.push(Task(title: "Test code", priority: 2))
queue2.push(Task(title: "Refactor code", priority: 2))

// Compare the task queues using the protocol extension method
let comparisonResult = queue1.compare(queue: queue2)

// Print the comparison result
switch comparisonResult {
case .orderedAscending:
    print("Queue 1 has fewer tasks than Queue 2")
case .orderedDescending:
    print("Queue 1 has more tasks than Queue 2")
case .orderedSame:
    print("Both queues have the same number of tasks")
}

// Output:
// Queue 1 has fewer tasks than Queue 2

Since queue1 has 2 tasks and queue2 has 3 tasks, the comparison result is .orderedAscending, indicating that queue1 has fewer tasks than queue2.

You can see see how we added the compare(queue:) method to all types conforming to the Queue protocol, without modifying their individual implementations. This kind of extension allows you to add functionality to protocols and their conforming types in a clean and modular way.

Swift Standard Library Extensions

Protocol extensions are extremely powerful when used with Swift's standard libraries. It enables you to add new functionalities to types conforming to standard protocols like Collection, Sequence, Equatable, etc.

For example, you can extend the Array type to provide custom logging capabilities:

extension Array where Element: CustomStringConvertible {
    func logElements() {
        for element in self {
            print(element)
        }
    }
}

let countries = ["United States", "Canada", "United Kingdom", "Australia", "Japan"]
countries.logElements()

// Output:
/*
 United States
 Canada
 United Kingdom
 Australia
 Japan
 */

In the above example, we've extended the Array collection to provide a custom method called logElements. This method can be called on any type that conforms to CustomStringConvertibleprotocol.

String Trimming for Collection:

extension Collection where Element: StringProtocol {
    func trimmedStrings() -> [String] {
        return self.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
    }
}

// Usage
let textArray = ["  Hello ", " Swift   ", "   Anytime  "]
let trimmedArray = textArray.trimmedStrings()
print(trimmedArray) // Print: ["Hello", "Swift", "Anytime"]

Here, we've extended the Collection protocol for elements that conform to StringProtocol (which includes String and Substring). We've added a method called trimmedStrings() that trims whitespace and newline characters from each element in the collection.

These examples illustrate how protocol extensions can be used to add useful and reusable functionality to existing types from the standard library, as well as to your custom types that conform to those protocols.

Where to go next?

Congratulations, you have mastered Protocol Extensions concept and this will let you make your codebase more readable and less redundant. You can checkout complete guide to Protocol Oriented Programming here or learn about Protocol Inheritance and Compositions in order to make full use of POP in Swift.