Browse Source

Restrictions: Significantly clean up code

main
Ave Ozkal 2 years ago
parent
commit
3b37a6c87a
3 changed files with 41 additions and 27 deletions
  1. +6
    -9
      cogs/logs.py
  2. +2
    -2
      cogs/mod.py
  3. +33
    -16
      helpers/restrictions.py

+ 6
- 9
cogs/logs.py View File

@ -1,6 +1,7 @@
import discord
import json
import config
from helpers.restrictions import get_user_restrictions
class Logs:
@ -43,15 +44,11 @@ class Logs:
f"🕓 Account age: {age}\n"\
f"🏷 __User ID__: {member.id}"
# Taken from kurisu source.
# Blame ihaveamac, not me.
with open("data/restrictions.json", "r") as f:
rsts = json.load(f)
if str(member.id) in rsts:
roles = []
for rst in rsts[str(member.id)]:
roles.append(discord.utils.get(member.guild.roles, id=rst))
await member.add_roles(*roles)
# Handles user restrictions
# Basically, gives back muted role to users that leave with it.
rsts = get_user_restrictions(member.id)
roles = [discord.utils.get(member.guild.roles, id=rst) for rst in rsts]
await member.add_roles(*roles)
# Real hell zone.
with open("data/userlog.json", "r") as f:


+ 2
- 2
cogs/mod.py View File

@ -57,7 +57,7 @@ class Mod:
log_channel = self.bot.get_channel(config.log_channel)
await log_channel.send(chan_message)
await ctx.send(f"{target.mention} can no longer speak.")
await add_restriction(target, config.mute_role)
add_restriction(target.id, config.mute_role)
@commands.guild_only()
@commands.check(check_if_staff)
@ -76,7 +76,7 @@ class Mod:
log_channel = self.bot.get_channel(config.log_channel)
await log_channel.send(chan_message)
await ctx.send(f"{target.mention} can now speak again.")
await remove_restriction(target, config.mute_role)
remove_restriction(target.id, config.mute_role)
@commands.guild_only()
@commands.bot_has_permissions(kick_members=True)


+ 33
- 16
helpers/restrictions.py View File

@ -1,25 +1,42 @@
import json
def add_restriction(self, member, rst):
# from kurisu source, credits go to ihaveamac
def get_restrictions():
with open("data/restrictions.json", "r") as f:
rsts = json.load(f)
if str(member.id) not in rsts:
rsts[str(member.id)] = []
if rst not in rsts[str(member.id)]:
rsts[str(member.id)].append(rst)
return json.load(f)
def set_restrictions(contents):
with open("data/restrictions.json", "w") as f:
json.dump(rsts, f)
f.write(contents)
def remove_restriction(self, member, rst):
# from kurisu source, credits go to ihaveamac
def get_user_restrictions(uid):
uid = str(uid)
with open("data/restrictions.json", "r") as f:
rsts = json.load(f)
if str(member.id) not in rsts:
rsts[str(member.id)] = []
if rst in rsts[str(member.id)]:
rsts[str(member.id)].remove(rst)
with open("data/restrictions.json", "w") as f:
json.dump(rsts, f)
if uid in rsts:
return rsts[uid]
return []
def add_restriction(uid, rst):
# mostly from kurisu source, credits go to ihaveamac
uid = str(uid)
rsts = get_restrictions()
if uid not in rsts:
rsts[uid] = []
if rst not in rsts[uid]:
rsts[uid].append(rst)
set_restrictions(json.dumps(rsts))
def remove_restriction(uid, rst):
# mostly from kurisu source, credits go to ihaveamac
uid = str(uid)
rsts = get_restrictions()
if uid not in rsts:
rsts[uid] = []
if rst in rsts[uid]:
rsts[uid].remove(rst)
set_restrictions(json.dumps(rsts))

Loading…
Cancel
Save