Source code for arize.users.types

"""Public type re-exports and SDK-facing role types for the users subdomain.

Domain role types (``PredefinedUserRole``, ``CustomUserRole``) subclass the
strict ``*Request`` schemas and are used as inputs to write methods. They are
also produced by ``User._coerce_role`` when deserializing responses.
"""

from datetime import datetime
from enum import Enum
from typing import Annotated, Literal

from pydantic import BaseModel, Field, field_validator

from arize._generated.api_client.models.create_user_request import (
    CreateUserRequest,
)
from arize._generated.api_client.models.custom_user_role_assignment import (
    CustomUserRoleAssignment,
)
from arize._generated.api_client.models.custom_user_role_assignment_request import (
    CustomUserRoleAssignmentRequest,
)
from arize._generated.api_client.models.invite_mode import InviteMode
from arize._generated.api_client.models.pagination_metadata import (
    PaginationMetadata,
)
from arize._generated.api_client.models.predefined_user_role_assignment import (
    PredefinedUserRoleAssignment,
)
from arize._generated.api_client.models.predefined_user_role_assignment_request import (
    PredefinedUserRoleAssignmentRequest,
)
from arize._generated.api_client.models.update_user_request import (
    UpdateUserRequest,
)
from arize._generated.api_client.models.user_role import UserRole
from arize._generated.api_client.models.user_role_assignment import (
    UserRoleAssignment,
)
from arize._generated.api_client.models.user_status import UserStatus




class PredefinedUserRole(PredefinedUserRoleAssignmentRequest):
    """A predefined account-level role assignment.

    The ``type`` discriminator is set to ``"PREDEFINED"`` automatically.

    Args:
        name: The predefined role name (``"ADMIN"``, ``"MEMBER"``, or
            ``"ANNOTATOR"``).
    """

    type: Literal["PREDEFINED"] = "PREDEFINED"  # type: ignore[assignment]

    def __str__(self) -> str:
        """Return the role name as a string."""
        return self.name.value





class CustomUserRole(CustomUserRoleAssignmentRequest):
    """A custom RBAC role assignment.

    The ``type`` discriminator is set to ``"CUSTOM"`` automatically.

    Args:
        id: The unique identifier of the custom RBAC role.
        name: Human-readable name of the custom role (returned in responses
            only; ignored on input).
    """

    # name is a response-only field not present in CustomUserRoleAssignmentRequest;
    # carried here so __str__ can display it when available.
    name: str | None = None
    type: Literal["CUSTOM"] = "CUSTOM"  # type: ignore[assignment]

    def __str__(self) -> str:
        """Return the role name if available, otherwise the role id."""
        return self.name if self.name is not None else self.id



_UserRoleField = Annotated[
    PredefinedUserRole | CustomUserRole,
    Field(discriminator="type"),
]




class User(BaseModel):
    """An account user with domain-typed role."""

    id: str
    name: str
    email: str
    created_at: datetime
    status: UserStatus
    role: _UserRoleField
    is_developer: bool

    @field_validator("role", mode="before")
    @classmethod
    def _coerce_role(cls, v: object) -> object:
        if isinstance(v, UserRoleAssignment):
            actual = v.actual_instance
            if isinstance(actual, PredefinedUserRoleAssignment):
                return PredefinedUserRole(name=actual.name)
            if isinstance(actual, CustomUserRoleAssignment):
                return CustomUserRole(id=actual.id, name=actual.name)
            raise TypeError(f"Unknown role type: {type(actual)!r}")
        return v





class CreateUserResponse(User):
    """Response from :meth:`UsersClient.create`.

    Extends :class:`User` with the invite mode used to create the account
    and, when ``invite_mode`` is ``"TEMPORARY_PASSWORD"``, the one-time
    password issued for the new account.
    """

    invite_mode: InviteMode
    # repr=False so logging/printing a CreateUserResponse (e.g. `logger.debug("%r", result)`)
    # doesn't leak the plaintext password, mirroring the generated model's use of SecretStr.
    temporary_password: str | None = Field(default=None, repr=False)





class ListUsersResponse(BaseModel):
    """Paginated list of users with domain-typed roles."""

    users: list[User]
    pagination: PaginationMetadata





class DeletionStatus(str, Enum):
    """Outcome of a single user deletion attempt."""

    DELETED = "deleted"
    FAILED = "failed"
    NOT_FOUND = "not_found"





class BulkUserDeletionResult(BaseModel):
    """Result of a single user deletion attempt.

    Attributes:
        user_id: ID of the user targeted for deletion. Empty when an email
            could not be resolved to a user (status ``"not_found"``).
        email: Email address the deletion was requested for, set only when
            the user was specified by email. ``None`` when specified by ID.
        status: Outcome of the deletion attempt.
        error: Error message when ``status`` is ``"failed"`` or
            ``"not_found"``.
    """

    user_id: str
    status: DeletionStatus
    email: str | None = None
    error: str | None = None





class BulkDeleteResponse(BaseModel):
    """Response from a bulk user delete operation.

    Attributes:
        results: Per-user deletion outcomes.
    """

    results: list[BulkUserDeletionResult]



__all__ = [
    "BulkDeleteResponse",
    "BulkUserDeletionResult",
    "CreateUserRequest",
    "CreateUserResponse",
    "CustomUserRole",
    "DeletionStatus",
    "InviteMode",
    "ListUsersResponse",
    "PredefinedUserRole",
    "UpdateUserRequest",
    "User",
    "UserRole",
    "UserStatus",
]