TeX-Bot-Py-V2/cogs/stats/counts.py at main · CSSUoB/TeX-Bot-Py-V2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""Contains methods relating to counting messages in channels."""

from typing import TYPE_CHECKING

import discord

from config import settings

if TYPE_CHECKING:
    from collections.abc import AsyncIterable, Mapping, Sequence


__all__: "Sequence[str]" = ("get_channel_message_counts", "get_server_message_counts")


async def get_channel_message_counts(channel: discord.TextChannel) -> "Mapping[str, int]":
    """
    Get the message counts for each role in the given channel.

    The message counts are stored in a mapping with the role name (prefixed by `@`) as the key
    and the number of messages sent by users with that role as the value.
    The mapping also includes a "Total" key for the total number of messages.
    """
    message_counts: dict[str, int] = {"Total": 0}

    role_name: str
    for role_name in settings["STATISTICS_ROLES"]:
        if discord.utils.get(channel.guild.roles, name=role_name):
            message_counts[f"@{role_name}"] = 0

    message_history_period: AsyncIterable[discord.Message] = channel.history(
        after=discord.utils.utcnow() - settings["STATISTICS_DAYS"]
    )
    message: discord.Message
    async for message in message_history_period:
        if message.author.bot:
            continue

        message_counts["Total"] += 1

        if isinstance(message.author, discord.User):
            continue

        author_role_names: set[str] = {
            author_role.name for author_role in message.author.roles
        }

        author_role_name: str
        for author_role_name in author_role_names:
            if f"@{author_role_name}" in message_counts:
                is_author_role_name: bool = author_role_name == "Committee"
                if is_author_role_name and "Committee-Elect" in author_role_names:
                    continue

                if author_role_name == "Guest" and "Member" in author_role_names:
                    continue

                message_counts[f"@{author_role_name}"] += 1

    return message_counts


async def get_server_message_counts(
    guild: discord.Guild, *, guest_role: discord.Role
) -> "Mapping[str, Mapping[str, int]]":
    """
    Get the message counts for each channel in the given server.

    The message counts are stored in a mapping. It contains a key "roles", which is
    a mapping of role names (prefixed by `@`) to the message counts
    for each role across the entire server.
    The mapping also contains a key "channels", which is a mapping with the channel
    name as a key and the number of messages sent in that channel as the value.
    The "roles" sub-mapping also includes a "Total" key for the total number of messages.
    """
    message_counts: dict[str, dict[str, int]] = {"roles": {"Total": 0}, "channels": {}}

    role_name: str
    for role_name in settings["STATISTICS_ROLES"]:
        if discord.utils.get(guild.roles, name=role_name):
            message_counts["roles"][f"@{role_name}"] = 0

    channel: discord.TextChannel
    for channel in guild.text_channels:
        member_has_access_to_channel: bool = channel.permissions_for(guest_role).is_superset(
            discord.Permissions(send_messages=True)
        )
        if not member_has_access_to_channel:
            continue

        message_counts["channels"][f"#{channel.name}"] = 0

        message_history_period: AsyncIterable[discord.Message] = channel.history(
            after=discord.utils.utcnow() - settings["STATISTICS_DAYS"]
        )
        message: discord.Message
        async for message in message_history_period:
            if message.author.bot:
                continue

            message_counts["channels"][f"#{channel.name}"] += 1
            message_counts["roles"]["Total"] += 1

            if isinstance(message.author, discord.User):
                continue

            author_role_names: set[str] = {
                author_role.name for author_role in message.author.roles
            }

            author_role_name: str
            for author_role_name in author_role_names:
                if f"@{author_role_name}" in message_counts["roles"]:
                    is_author_role_committee: bool = author_role_name == "Committee"
                    if is_author_role_committee and "Committee-Elect" in author_role_names:
                        continue

                    if author_role_name == "Guest" and "Member" in author_role_names:
                        continue

                    message_counts["roles"][f"@{author_role_name}"] += 1

    return message_counts