-
Notifications
You must be signed in to change notification settings - Fork 5
/
timezone_bot.py
131 lines (93 loc) · 3.58 KB
/
timezone_bot.py
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
124
125
126
127
128
129
130
131
# Original author: https://github.com/JakeStanger/Discord-Friend-Timezones/tree/6eb9537bc9c9e434549c001c1859635c82333b66
import os
from datetime import datetime
import discord
import pytz
from discord.ext import commands
from fuzzywuzzy import process
from sqlalchemy import Column, String
from sqlalchemy.ext.declarative import declarative_base
import utilities
session = utilities.get_timezone_session()
Base = declarative_base()
class UserZone(Base):
__tablename__ = 'UserZone'
id = Column(String(18), primary_key=True)
zone = Column(String(64))
async def get_or_create(session, model, **kwargs):
instance = session.query(model).filter_by(**kwargs).first()
if instance:
return instance
else:
instance = model(**kwargs)
session.add(instance)
session.commit()
return instance
description = f"""
A bot that manages your timezone
Typical usage:
1. Know the 2-letter code for your country.
You may get it from https://gist.github.com/Zackhardtoname/c2a01465ac25eec17038c6e871be17ef
2. Get available timezones in your country with .tzlist your_code (example: .tzlist US)
3. Set it with .tzset one_of_the_timezones (example: .tzset America/New_York)
(4. Check it with .tzget and .time)
"""
timezone_bot = commands.Bot(command_prefix=utilities.config['timezone_prefix'], description=description)
@timezone_bot.command(name='tzset')
async def set_zone(ctx, *, timezone):
if timezone in pytz.all_timezones:
zone = timezone
else:
zone = process.extractOne(timezone, pytz.all_timezones)[0]
user_zone = await get_or_create(session, UserZone, id=ctx.author.id)
user_zone.zone = zone
session.commit()
await ctx.send("Set your time zone to **%s**" % zone)
async def query_zone(user: discord.Member):
zone = session.query(UserZone).filter_by(id=user.id).first()
if zone:
return zone.zone
else:
return 'Not set'
async def get_zone_time(zone: str, full=False):
if zone == 'Not set':
return zone
tz = pytz.timezone(zone)
res = tz.normalize(datetime.utcnow().replace(tzinfo=pytz.utc).astimezone(tz))
if not full:
res = res.strftime('%H:%M')
return res
@set_zone.error
async def info_error(ctx, error):
if isinstance(error, commands.MissingRequiredArgument):
await ctx.send('Please specify a timezone')
else:
print(error)
@timezone_bot.command(name='tzget')
async def get_zone(ctx, users: commands.Greedy[discord.Member]):
if len(users) == 0:
users = [ctx.author]
await ctx.send('\n'.join([
'%s: **%s**' % (user.nick or user.name,
await query_zone(user))
for user in set(users)
]))
@timezone_bot.command(name='time')
async def get_time(ctx, users: commands.Greedy[discord.Member]):
if len(users) == 0:
users = [ctx.author]
await ctx.send('\n'.join([
'%s: **%s**' % (user.nick or user.name,
await get_zone_time(await query_zone(user)))
for user in set(users)
]))
@timezone_bot.command(name='tzlist', help='Get a list of available timezones')
async def get_tzlist(ctx, country=None):
if country != None and len(country) == 2:
await ctx.send(
f'Available timezones for {pytz.country_names[country.upper()]} are:\n{", ".join(pytz.country_timezones[country.upper()])}')
else:
await ctx.send(
'Please specify this command with a two letter country code. [List available here](https://gist.github.com/pamelafox/986163)')
if __name__ == '__main__':
timezone_bot.run(os.getenv("timezone_token"))