From 990d8cfc98ec71a31ea3ef659cb109547f987e31 Mon Sep 17 00:00:00 2001
From: ks129 <45097959+ks129@users.noreply.github.com>
Date: Sat, 10 Oct 2020 19:56:54 +0300
Subject: Include voice ban status on metricity data

---
 pydis_site/apps/api/viewsets/bot/user.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

(limited to 'pydis_site')

diff --git a/pydis_site/apps/api/viewsets/bot/user.py b/pydis_site/apps/api/viewsets/bot/user.py
index 3ab71186..19f55555 100644
--- a/pydis_site/apps/api/viewsets/bot/user.py
+++ b/pydis_site/apps/api/viewsets/bot/user.py
@@ -1,6 +1,7 @@
 import typing
 from collections import OrderedDict
 
+from django.core.exceptions import MultipleObjectsReturned, ObjectDoesNotExist
 from rest_framework import status
 from rest_framework.decorators import action
 from rest_framework.pagination import PageNumberPagination
@@ -9,6 +10,7 @@ from rest_framework.response import Response
 from rest_framework.serializers import ModelSerializer
 from rest_framework.viewsets import ModelViewSet
 
+from pydis_site.apps.api.models.bot.infraction import Infraction
 from pydis_site.apps.api.models.bot.metricity import Metricity, NotFound
 from pydis_site.apps.api.models.bot.user import User
 from pydis_site.apps.api.serializers import UserSerializer
@@ -240,10 +242,21 @@ class UserViewSet(ModelViewSet):
     def metricity_data(self, request: Request, pk: str = None) -> Response:
         """Request handler for metricity_data endpoint."""
         user = self.get_object()
+
+        try:
+            Infraction.objects.get(user__id=user.id, active=True, type="voice_ban")
+        except ObjectDoesNotExist:
+            voice_banned = False
+        except MultipleObjectReturned:
+            voice_banned = True
+        else:
+            voice_banned = True
+
         with Metricity() as metricity:
             try:
                 data = metricity.user(user.id)
                 data["total_messages"] = metricity.total_messages(user.id)
+                data["voice_banned"] = voice_banned
                 return Response(data, status=status.HTTP_200_OK)
             except NotFound:
                 return Response(dict(detail="User not found in metricity"),
-- 
cgit v1.2.3