import jinja2
from jinja2.ext import Extension

from .wagtailuserbar import wagtailuserbar


class WagtailUserbarExtension(Extension):
    def __init__(self, environment):
        super(WagtailUserbarExtension, self).__init__(environment)

        self.environment.globals.update({
            'wagtailuserbar': jinja2.contextfunction(wagtailuserbar),
        })


# Nicer import names
userbar = WagtailUserbarExtension

Add `wagtailuserbar()` Jinja function