Registering custom context initializers with behat 3

behat.yml

default:
    extensions:
        custom_extension.php: ~

extension.php

<?php

use Symfony\Component\Config\Definition\Builder\ArrayNodeDefinition;
use Symfony\Component\DependencyInjection\ContainerBuilder;
use Behat\Testwork\ServiceContainer\ExtensionManager;
use Behat\Testwork\ServiceContainer\Extension;
use Behat\Behat\Context\ServiceContainer\ContextExtension;
use Behat\Behat\Context\Initializer\ContextInitializer;
use Behat\Behat\Context\Context;

class CustomInitializer implements ContextInitializer
{
    public function supportsContext(Context $context)
    {
        return true;
    }

    public function initializeContext(Context $context)
    {
        // do the stuff with the $context
    }
}

class CustomExtension implements Extension
{
    public function getConfigKey()
    {
        return 'custom_extension';
    }

    public function load(ContainerBuilder $container, array $config)
    {
        $definition = $container->register('custom_initializer', 'CustomInitializer');
        $definition->addTag(ContextExtension::INITIALIZER_TAG);
    }

    public function configure(ArrayNodeDefinition $builder) {}
    public function initialize(ExtensionManager $extensionManager) {}
    public function process(ContainerBuilder $container) {}
}

return new CustomExtension;