from wagtail.users.views.users import UserViewSet as WagtailUserViewSet

from .forms import WagtailUserCreateForm, WagtailUserEditForm


class UserViewSet(WagtailUserViewSet):
    def get_form_class(self, for_update=False):
        """
        If for_update=True, this is editing an existing user,
        else it's for creating a new user.
        """
        if for_update:
            return WagtailUserEditForm
        return WagtailUserCreateForm
