from dataclasses import dataclass
from .base import CallType, TSEndpoint
from .decorators import authenticate
from .mocks import traits_data
from .constants import BATCH, TS, TRAITS, USERS
[docs]
@dataclass
class UserTrait:
"""Model that describes the payload for the set_user_traits endpoint."""
trait_id: str
"""The id of the trait being set for the user"""
trait_value: str | bool | float | int
"""The value of the trait. The data type depends on the trait being set."""
[docs]
@dataclass
class TraitBatch:
"""Model that describes the payload for the set_trait_batch endpoint."""
user_id: int
"""The user id the trait is being set for"""
trait_value: str | bool | float | int
"""The value of the trait. The data type depends on the trait being set."""
[docs]
class Traits(TSEndpoint):
[docs]
@authenticate
async def set_trait(self, trait_id: str, user_id: int) -> bool:
"""Deprecated. Please use set_user_traits instead."""
response = await self._send_request(
f"{TS}/{TRAITS}/{trait_id}/{USERS}/{user_id}",
CallType.POST,
mock_response=self.mock_set_trait,
)
self.logger.warning(
"The `set_trait` function is deprecated. Please use `set_user_traits` instead"
)
return response.ok
@property
def mock_set_trait(self) -> dict:
return {}
def _process_traits_for_trait_store(self, trait_objects: list) -> dict:
"""
Convert a list of trait objects to a dictionary format for the trait store API.
:param trait_objects: List of UserTrait or TraitBatch objects to process
:return: Dictionary mapping IDs to trait values (trait_id -> trait_value for UserTrait, user_id -> trait_value for TraitBatch)
"""
values = {}
for trait_object in trait_objects:
if isinstance(trait_object, UserTrait):
values[trait_object.trait_id] = trait_object.trait_value
elif isinstance(trait_object, TraitBatch):
values[trait_object.user_id] = trait_object.trait_value
return values
[docs]
@authenticate
async def set_user_traits(self, user_id: int, user_traits: list[UserTrait]) -> dict:
"""
Set multiple traits for a specific user in batch.
:param user_id: The ID of the user to set traits for
:param user_traits: List of UserTrait objects containing trait_id and trait_value pairs
:return: Response data indicating the number of success, or error information
"""
values = self._process_traits_for_trait_store(user_traits)
response = await self._send_request(
f"{TS}/{TRAITS}/{USERS}/{user_id}/{BATCH}",
CallType.POST,
mock_response=self.mock_set_user_traits,
json=values,
)
return response.json()
@property
def mock_set_user_traits(self) -> dict:
return traits_data.set_user_traits_mock()
[docs]
@authenticate
async def set_trait_batch(
self, trait_id: str, trait_batch: list[TraitBatch]
) -> dict:
"""
Set a specific trait for multiple users in batch.
:param trait_id: The ID of the trait to set
:param trait_batch: List of TraitBatch objects containing user_id and trait_value pairs
:return: Response data indicating the number of successfully set traits, and any error information
"""
values = self._process_traits_for_trait_store(trait_batch)
response = await self._send_request(
f"{TS}/{TRAITS}/{trait_id}/{USERS}/{BATCH}",
CallType.POST,
mock_response=self.mock_set_trait_batch,
json=values,
)
return response.json()
@property
def mock_set_trait_batch(self) -> dict:
return traits_data.set_trait_batch_mock()