71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
"""Data models for chat application."""
|
|
|
|
import asyncio
|
|
from typing import Set
|
|
|
|
|
|
class Channel:
|
|
"""Represents a chat channel or group."""
|
|
def __init__(self, name: str, creator: str, is_private: bool = False):
|
|
self.name = name
|
|
self.creator = creator
|
|
self.is_private = is_private
|
|
self.members: Set[str] = {creator}
|
|
self.messages = []
|
|
self.invites: Set[str] = set() # Pending invites
|
|
self.max_messages = 500
|
|
|
|
def add_member(self, username: str):
|
|
self.members.add(username)
|
|
if username in self.invites:
|
|
self.invites.remove(username)
|
|
|
|
def remove_member(self, username: str):
|
|
if username in self.members:
|
|
self.members.remove(username)
|
|
|
|
def add_message(self, msg_data: dict):
|
|
self.messages.append(msg_data)
|
|
if len(self.messages) > self.max_messages:
|
|
self.messages.pop(0)
|
|
|
|
def has_member(self, username: str) -> bool:
|
|
return username in self.members
|
|
|
|
def is_invited(self, username: str) -> bool:
|
|
return username in self.invites
|
|
|
|
def invite(self, username: str):
|
|
if username not in self.members:
|
|
self.invites.add(username)
|
|
|
|
|
|
class DirectMessage:
|
|
"""Represents a DM conversation between two users."""
|
|
def __init__(self, user1: str, user2: str):
|
|
self.users = frozenset([user1, user2])
|
|
self.messages = []
|
|
self.max_messages = 500
|
|
|
|
def add_message(self, msg_data: dict):
|
|
self.messages.append(msg_data)
|
|
if len(self.messages) > self.max_messages:
|
|
self.messages.pop(0)
|
|
|
|
def get_other_user(self, username: str) -> str:
|
|
users_list = list(self.users)
|
|
return users_list[0] if users_list[1] == username else users_list[1]
|
|
|
|
|
|
class UserConnection:
|
|
"""Represents a connected user with a message queue."""
|
|
def __init__(self, username: str, message_queue: asyncio.Queue):
|
|
self.username = username
|
|
self.message_queue = message_queue
|
|
|
|
def send_message(self, msg_data: dict):
|
|
try:
|
|
self.message_queue.put_nowait(msg_data)
|
|
except:
|
|
pass
|