Add missing code

[yoelcortes]

There was code missing in the previous pull request (which included the doctest) that went unnoticed.