Files
django-discord/chat/consumers.py
T

160 lines
4.9 KiB
Python

import json
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.utils import timezone
class ChannelConsumer(AsyncWebsocketConsumer):
async def connect(self):
self.channel_id = self.scope['url_route']['kwargs']['channel_id']
self.group_name = f'channel_{self.channel_id}'
self.user = self.scope['user']
if not self.user.is_authenticated:
await self.close()
return
await self.channel_layer.group_add(self.group_name, self.channel_name)
await self.set_online(True)
await self.accept()
async def disconnect(self, close_code):
await self.channel_layer.group_discard(self.group_name, self.channel_name)
await self.set_online(False)
async def receive(self, text_data):
data = json.loads(text_data)
content = data.get('message', '').strip()
if not content:
return
msg = await self.save_message(content)
await self.channel_layer.group_send(
self.group_name,
{
'type': 'chat_message',
'message': content,
'sender': self.user.username,
'sender_id': self.user.id,
'message_id': msg.id,
'timestamp': msg.timestamp.strftime('%H:%M'),
'avatar': await self.get_avatar(),
'role': await self.get_role(),
}
)
async def chat_message(self, event):
await self.send(text_data=json.dumps(event))
@database_sync_to_async
def save_message(self, content):
from chat.models import Message, Channel
channel = Channel.objects.get(id=self.channel_id)
return Message.objects.create(channel=channel, sender=self.user, content=content)
@database_sync_to_async
def set_online(self, status):
try:
self.user.profile.online = status
self.user.profile.save(update_fields=['online'])
except Exception:
pass
@database_sync_to_async
def get_avatar(self):
try:
av = self.user.profile.avatar
return av.url if av else None
except Exception:
return None
@database_sync_to_async
def get_role(self):
try:
return self.user.profile.role
except Exception:
return 'user'
class DMConsumer(AsyncWebsocketConsumer):
async def connect(self):
self.user = self.scope['user']
self.other_username = self.scope['url_route']['kwargs']['username']
if not self.user.is_authenticated:
await self.close()
return
other_id = await self.get_other_id()
if other_id is None:
await self.close()
return
ids = sorted([self.user.id, other_id])
self.group_name = f'dm_{ids[0]}_{ids[1]}'
await self.channel_layer.group_add(self.group_name, self.channel_name)
await self.set_online(True)
await self.accept()
async def disconnect(self, close_code):
await self.channel_layer.group_discard(self.group_name, self.channel_name)
await self.set_online(False)
async def receive(self, text_data):
data = json.loads(text_data)
content = data.get('message', '').strip()
if not content:
return
msg = await self.save_dm(content)
await self.channel_layer.group_send(
self.group_name,
{
'type': 'dm_message',
'message': content,
'sender': self.user.username,
'sender_id': self.user.id,
'message_id': msg.id,
'timestamp': msg.timestamp.strftime('%H:%M'),
'avatar': await self.get_avatar(),
}
)
async def dm_message(self, event):
await self.send(text_data=json.dumps(event))
@database_sync_to_async
def get_other_id(self):
from django.contrib.auth.models import User
try:
return User.objects.get(username=self.other_username).id
except User.DoesNotExist:
return None
@database_sync_to_async
def save_dm(self, content):
from django.contrib.auth.models import User
from chat.models import DirectMessage
other = User.objects.get(username=self.other_username)
return DirectMessage.objects.create(sender=self.user, receiver=other, content=content)
@database_sync_to_async
def set_online(self, status):
try:
self.user.profile.online = status
self.user.profile.save(update_fields=['online'])
except Exception:
pass
@database_sync_to_async
def get_avatar(self):
try:
av = self.user.profile.avatar
return av.url if av else None
except Exception:
return None