Fresh Beginning

Implementing custom NotificationCenter in Swift

Jayesh Kawli -

NSNotificationCenter is a classic API on iOS. It's used so that classes can register for notifications by name and execute the function/block whenever the notification with same name is fired somewhere else in the project.

Today, we will see how to implement NSNotificationCenter API in Swift from scratch. Here are things that we want this implementation to do,

  1. Add class as an observer for notification
  2. Post a notification with a name
  3. Remove class as an observer

Since any class in the project can register for notifications, our notification center will be declared as a singleton. Anyone who wish to register/listen for/remove the notification will use centralized singleton object.

Data structure for notification center storage

When we store the notifications, there are two levels of storage. First one is class which is an observer and second the name of notification and the code we want to execute.

  1. One class can register to multiple notifications with different names
  2. One notification can execute different source code

In order to maintain such hierarchy, we will use the nested dictionary with an array container.

  1. First level of dictionary will store class name
  2. Second level will store all the keys representing notification names
  3. Third level will store all the closure/code to execute associated with that class name - notification name pair

So finally that dictionary will look like this in Swift. Please note how we're using closures as code execution blocks.

First parameter in the closure indicates the name of notification and second represents any data you want to pass through the notification.


// Storing first string as a class name nested with the notification name
// Which is further nested with closure to execute when notification is posted

private var notificationsStorage: [String: [String: [(String, Any) -> Void]]]

Custom NotificationCenter Singleton class

As mentioned earlier, our custom notification center is going to be a singleton only one instance will be used by anyone using this API in the app.


final class MyNotificationCenter {
    static let shared = MyNotificationCenter()
    
    // Storing first string as a class name nested with notification name
    // Which is further nested with closure to execute when notification is posted
    private var notificationsStorage: [String: [String: [(String, Any) -> Void]]]

    private init() {
        notificationsStorage = [:]
    }
  ...
  ..
  .
}

Note how we're initializing notificationsStorage dictionary which will store notifications data.

Every time we need to use MyNotificationCenter, we will access it with MyNotificationCenter.shared variable.

Adding class as an observer

The first part of our API is to let us add an instance of class as an observer. One class may be observer to multiple notifications. There could also be the case where multiple classes can register to same notification but may execute different actions.

Our addObserver API will thus take 3 parameters,

  1. Class - An instance of class which acts as an observer
  2. Name - Name of the notification
  3. Closure - This is the block of code to execute after notification is fired. Closure may contain arbitrary code, but it will always get called with two parameters, notification name and object - The data that we want to pass with the notification

Below is the full method implementation,


func addObserver(_ _class: Any, name: String, closure: @escaping (String, Any) -> Void) {
    // If the instance is not of a class, ignore.
    // We only support adding observer to class instances
    guard let inputClass = type(of: _class) as? AnyClass else {
        return
    }
    
    // Convert instance to class name
    let className = String(describing: inputClass)
    
    // If we already have an entry for that class and notification name,
    // append the input closure to existing array in the dictionary
    if notificationsStorage[className] != nil && notificationsStorage[className]?[name] != nil {
        notificationsStorage[className]?[name]?.append(closure)
    } else if notificationsStorage[className] != nil {
    	// If class name and notification name pair do not exist, create a new array
    	if notificationsStorage[className]?[name] == nil {
 		   notificationsStorage[className]?[name] = [closure]
		} else {
            // If pair exists, append the new closure to existing class name/notification name pair
		    notificationsStorage[className]?[name]?.append(closure)
		}
    } else {
        // If the class and notification name pair does not exist,
        // Add it to the dictionary with closure
        notificationsStorage[className] = [name: [closure]]
	}
}
Thanks to Sagar Daundkar for pointing out the bug in above code which prevented it from adding multiple notifications on the same observer

Posting notification

What happens when we post the notification? All the observers observing notification with that name get called and corresponding blocks get called. When someone posts the notification with that name and object parameter, we will go through the notificationsStorage container, find all notifications with that name and fire them one by one.

If we cannot find any notification block by that name, we will simply throw the notificationNotFound exception.


enum NotificationCenterError: Error {
    case notificationNotFound
}

....
...
..

func postNotification(_ name: String, object: Any) throws {
    
    var atLeastOneNotificationFound = false
    for (_, notificationData) in notificationsStorage {
        // Go through all the notifications and do the
        // Matching with name
        for (notificationName, closures) in notificationData {
            // Check if the notification name matches
            guard notificationName == name else { continue }
            for closure in closures {
                atLeastOneNotificationFound = true
                closure(name, object)
            }
        }
    }
    
    if !atLeastOneNotificationFound {
        throw NotificationCenterError.notificationNotFound
    }
}

Removing an observer

Lastly, we may want to remove a class as an observer to prevent memory leaks. When class instance is removed as an observer, it won't received any kind of notifications. In our case, this is fairly simple. We have notificationsStorage as a container which stores all notifications data keyed at top level by class name. Thus when the observer wishes to be removed, it will simply pass its instance to removeObserver method.


func removeObserver(_ _class: Any) throws {
    
    guard let inputClass = type(of: _class) as? AnyClass else {
        return
    }
    
    let className = String(describing: inputClass)
    
    guard notificationsStorage[className] != nil else { throw NotificationCenterError.notificationNotFound }
    notificationsStorage.removeValue(forKey: className)
}

If the observer is not of type class, we will return and result in no-op. If no notifications are found keyed by that class name, we will throw notificationNotFound exception.

Eventually we will remove the value keyed by input class name

How to use an API?

Here is a small demo on how to use the API


class Test {

    func foo() {
        let notificationCenter = MyNotificationCenter.shared
        
        notificationCenter.addObserver(self, name: "one") { _, _ in
            print("one")
        }
        
        notificationCenter.addObserver(self, name: "one") { _, _ in
            print("one one")
        }
        
        notificationCenter.addObserver(self, name: "one") { _, _ in
            print("one one one")
        }
        
        notificationCenter.addObserver(self, name: "two") { _, _ in
            print("Two")
        }
        
        notificationCenter.addObserver(self, name: "three") { _, _ in
            print("Three")
        }
        
        _ = try? notificationCenter.postNotification("one", object: "None")
        _ = try? notificationCenter.postNotification("two", object: "None")
        _ = try? notificationCenter.postNotification("three", object: "None")
    }


}

This will print,

one
one one
one one one
Two
Three

The full source code of MyNotificationCenter can be found in this gist.

Hope this was useful. I am always looking for a feedback so if you have any comments or suggestions to improve this script, would love to hear from you!