mizan/examples/django-react-site/backend/testapp/clients.py

from django import forms
from django.http import HttpRequest
from pydantic import BaseModel

from mizan.client import ServerFunction, client
from mizan.channels import Channel
from mizan.setup import register, register_form, register_as
from mizan.channels import register as register_channel
from mizan.forms import mizanFormMixin, mizanFormMeta
from mizan.jwt import jwt_obtain, jwt_refresh


# =============================================================================
# Server Functions
# =============================================================================


class EchoOutput(BaseModel):
    message: str


@client(websocket=True)
def echo(request: HttpRequest, text: str) -> EchoOutput:
    return EchoOutput(message=text)


register(echo, "echo")


class AddOutput(BaseModel):
    result: int


@client(websocket=True)
def add(request: HttpRequest, a: int, b: int) -> AddOutput:
    return AddOutput(result=a + b)


register(add, "add")


class WhoamiOutput(BaseModel):
    user_id: int | None
    email: str
    is_staff: bool


@client(auth=True)
def whoami(request: HttpRequest) -> WhoamiOutput:
    return WhoamiOutput(
        user_id=getattr(request.user, "id", None),
        email=getattr(request.user, "email", ""),
        is_staff=getattr(request.user, "is_staff", False),
    )


register(whoami, "whoami")


@client
def http_only_echo(request: HttpRequest, text: str) -> EchoOutput:
    return EchoOutput(message=text)


register(http_only_echo, "http_only_echo")


# =============================================================================
# Forms
# =============================================================================


class LoginForm(forms.Form):
    login = forms.CharField(max_length=150, label="Login")
    password = forms.CharField(widget=forms.PasswordInput, label="Password")


def handle_login(request, form):
    """Authenticate the credentials and open a session."""
    from django.contrib.auth import authenticate, login

    user = authenticate(
        request,
        username=form.cleaned_data["login"],
        password=form.cleaned_data["password"],
    )
    if user is not None:
        login(request, user)
        return {"success": True}
    form.add_error(None, "Invalid login credentials.")
    return None  # Signals validation failure


register_form(LoginForm, "login", submit_handler=handle_login)


class SignupForm(forms.Form):
    email = forms.EmailField(label="Email")
    password1 = forms.CharField(widget=forms.PasswordInput, label="Password")


def handle_signup(request, form):
    """Create the user, surfacing any creation failure as a form error."""
    from django.contrib.auth import get_user_model

    User = get_user_model()
    try:
        user = User.objects.create_user(
            email=form.cleaned_data["email"],
            password=form.cleaned_data["password1"],
        )
    except Exception as e:
        form.add_error(None, str(e))
        return None

    return {"success": True, "data": {"user_id": user.pk}}


register_form(SignupForm, "signup", submit_handler=handle_signup)


class AddEmailForm(forms.Form):
    email = forms.EmailField(label="Email address")


register_form(AddEmailForm, "add_email")


# =============================================================================
# Channels
# =============================================================================


class ChatChannel(Channel):
    class Params(BaseModel):
        room: str

    class ClientMessage(BaseModel):
        text: str

    class ServerMessage(BaseModel):
        text: str

    def authorize(self, params=None):
        return bool(params and params.room)

    def group(self, params=None):
        return f"chat_{params.room}"

    def receive(self, params, msg):
        return self.ServerMessage(text=msg.text)


register_channel(ChatChannel, "chat")


class NotificationsChannel(Channel):
    class ServerMessage(BaseModel):
        text: str

    def authorize(self, params=None):
        return self.user.is_authenticated

    def group(self, params=None):
        return "notifications_global"


register_channel(NotificationsChannel, "notifications")


class PresenceChannel(Channel):
    class ServerMessage(BaseModel):
        value: int

    def authorize(self, params=None):
        return self.user.is_authenticated

    def group(self, params=None):
        return "presence_global"


register_channel(PresenceChannel, "presence")


# =============================================================================
# Auth Variations
# =============================================================================


# --- Staff-only ---
@client(auth="staff")
def staff_only(request: HttpRequest) -> EchoOutput:
    return EchoOutput(message=f"staff:{request.user.email}")


register(staff_only, "staff_only")


# --- Superuser-only ---
@client(auth="superuser")
def superuser_only(request: HttpRequest) -> EchoOutput:
    return EchoOutput(message=f"superuser:{request.user.email}")


register(superuser_only, "superuser_only")


# --- Callable auth ---
def check_verified_email(request):
    if not request.user.is_authenticated:
        return False
    return getattr(request.user, "email", "").endswith("@verified.com")


@client(auth=check_verified_email)
def verified_only(request: HttpRequest) -> EchoOutput:
    return EchoOutput(message="verified")


register(verified_only, "verified_only")


# =============================================================================
# Context Functions
# =============================================================================


class CurrentUserOutput(BaseModel):
    authenticated: bool
    email: str
    is_staff: bool


@client(context="global")
def current_user(request: HttpRequest) -> CurrentUserOutput:
    if request.user.is_authenticated:
        return CurrentUserOutput(
            authenticated=True,
            email=request.user.email,
            is_staff=request.user.is_staff,
        )
    return CurrentUserOutput(authenticated=False, email="", is_staff=False)


register(current_user, "current_user")


class GreetOutput(BaseModel):
    greeting: str


@client(context="local")
def greet(request: HttpRequest, name: str) -> GreetOutput:
    return GreetOutput(greeting=f"Hello, {name}!")


register(greet, "greet")


# =============================================================================
# Class-based ServerFunction
# =============================================================================


class MultiplyInput(BaseModel):
    x: int
    y: int


class MultiplyOutput(BaseModel):
    product: int


@register_as("multiply")
class Multiply(ServerFunction):
    Input = MultiplyInput
    Output = MultiplyOutput

    def call(self, input: MultiplyInput) -> MultiplyOutput:
        return MultiplyOutput(product=input.x * input.y)


# =============================================================================
# Functions carrying error branches
# =============================================================================


_ECHO_TRANSFORMS = {"upper": str.upper, "lower": str.lower}


@client
def echo_transform(request: HttpRequest, text: str, mode: str) -> EchoOutput:
    transform = _ECHO_TRANSFORMS.get(mode)
    if transform is None:
        raise NotImplementedError(f"echo mode {mode!r} has no transform")
    return EchoOutput(message=transform(text))


register(echo_transform, "echo_transform")


class DivideOutput(BaseModel):
    quotient: float


@client
def divide(request: HttpRequest, numerator: int, denominator: int) -> DivideOutput:
    return DivideOutput(quotient=numerator / denominator)


register(divide, "divide")


@client
def permission_check_fn(request: HttpRequest, secret: str) -> EchoOutput:
    if secret != "open-sesame":
        raise PermissionError("Wrong secret")
    return EchoOutput(message="access granted")


register(permission_check_fn, "permission_check_fn")


# =============================================================================
# WebSocket + Auth Function
# =============================================================================


@client(websocket=True, auth=True)
def ws_whoami(request: HttpRequest) -> WhoamiOutput:
    return WhoamiOutput(
        user_id=getattr(request.user, "id", None),
        email=getattr(request.user, "email", ""),
        is_staff=getattr(request.user, "is_staff", False),
    )


register(ws_whoami, "ws_whoami")


# =============================================================================
# mizanFormMixin Forms
# =============================================================================


class ContactForm(mizanFormMixin, forms.Form):
    mizan = mizanFormMeta(
        name="contact",
        title="Contact Us",
        subtitle="We'd love to hear from you",
        submit_label="Send Message",
        live_validation=True,
        live_form_errors=False,
    )

    name = forms.CharField(max_length=100, label="Your Name")
    email = forms.EmailField(label="Email Address")
    message = forms.CharField(widget=forms.Textarea, label="Message")

    def on_submit_success(self, request):
        return {"received": True, "from": self.cleaned_data["email"]}


# =============================================================================
# Formset-enabled Form
# =============================================================================


class ItemForm(mizanFormMixin, forms.Form):
    mizan = mizanFormMeta(
        name="item",
        title="Items",
        submit_label="Save Items",
        enable_formset=True,
    )

    label = forms.CharField(max_length=50, label="Item Label")
    quantity = forms.IntegerField(min_value=1, label="Quantity")

    def on_submit_success(self, request):
        return {
            "label": self.cleaned_data["label"],
            "qty": self.cleaned_data["quantity"],
        }


# =============================================================================
# Staff-gated Channel
# =============================================================================


class PrivateChannel(Channel):
    class ServerMessage(BaseModel):
        text: str

    def authorize(self, params=None):
        return self.user.is_staff

    def group(self, params=None):
        return "private_global"


register_channel(PrivateChannel, "private")


# =============================================================================
# JWT Function Registration
# =============================================================================


register(jwt_obtain, "jwt_obtain")
register(jwt_refresh, "jwt_refresh")