Writer Monad and the Tracing Robot

After that small discussion on Y Combinator, we are back to monads! In the previous article on monads, we saw how we can use the the Result monad to chain a sequence of API calls together while automatically handling any failures along the way.

In this article we will learn about a new monad – the Writer monad. But first, let us take a look at the problem we are going to apply it to.

The Tracing Robot

We once again return to the robot problem from before. Here is the problem statement again:

• Imagine a robot that is at position (0, 0) and facing North.
• It can take 3 commands: Move forward, turn left or turn right.
• Turning left or right changes the orientation of the robot by 90 degrees without changing the position. For example, turn left will make the robot remain at (0, 0) but it will face West now
• Moving forward will move the robot 1 step in the direction faced, so if it is at (0, 0) and facing West, after moving forwards it will be at (-1, 0) and still facing West

Additionally, this time we want the robot to trace each step that it enters. The log should be something like this:

• Robot moved to (0, 1)
• Robot moved to (-1, 1)
• Robot moved to (-2, 1)
• ... and so on

The Writer Monad

This time apart from the RobotState we also need to keep track of the logs. These logs will be the additional context for our monad. This type of monad is called the Writer monad. (Remember, a monad is nothing but wrapping a value along with some additional context)

We can represent the monad like this

class Writer:
    def __init__(self, value, logs):
        self.value = value
        self.logs = logs

As before, we need a map method that will take an ordinary function and apply it to the value.

    def map(self, fn):
        return Writer(fn(self.value), self.logs)

Remember, map is used when we want to apply a normal function that just transforms the value without adding any new logs. So all we need to do is to unwrap the value, apply the function on it and wrap it again with the same logs.

We also need a flatmap method. This method will be used when we want to apply a function that itself returns some logs. In this case, we need to merge the existing logs in the monad with the new logs generated by the function.

    def flatmap(self, fn):
        output = fn(self.value)
        return Writer(output.value, self.logs + output.logs)

With that, our monad is complete.

Here is the full implementation.

class Writer:
    def __init__(self, value, logs):
        self.value = value
        self.logs = logs

    def map(self, fn):
        return Writer(fn(self.value), self.logs)

    def flatmap(self, fn):
        output = fn(self.value)
        return Writer(output.value, self.logs + output.logs)

Applying Writer Monad to the Robot

With the Writer monad now defined, we can apply it to the robot. The code for the left and right functions remain the same. Check the original robot kata for the implementation. The only change is in the move function. This function should not only calculate the new RobotState, but also return the appropriate log message along with it.

def log_position(robot: RobotState):
    (x, y), _ = robot
    return Writer(robot, [f"Robot moved to ({x}, {y})"])

move = compose(move, log_position)

With that single change, we can now simply compose the functions as usual.

start = make_robot((0, 0), "North")
end = (move(start)
        .map(left)
        .flatmap(move)
        .flatmap(move)
        .map(left)
        .flatmap(move)
        .flatmap(move)
        .flatmap(move)
        .map(right)
        .map(right)
        .map(right)
        .flatmap(move))

print(end.value)
print(end.logs)

At the end of the composition we can take out the final state as well as the list of logs generated by the sequence of operations. This is the output

((-1, -2), 'East')
['Robot moved to (0, 1)', 'Robot moved to (-1, 1)', 'Robot moved to (-2, 1)', 'Robot moved to (-2, 0)', 'Robot moved to (-2, -1)', 'Robot moved to (-2, -2)'
, 'Robot moved to (-1, -2)']

Isn't it great? We don't need to bother that the move function now generates logs, we just compose the functions as usual. The Writer monad takes care of all the logic handling logs and we don't need to worry about it.