henine
/
wubloader
mirror of https://github.com/ekimekim/wubloader
1
0
Fork 0
Code Issues Releases Wiki Activity

Combine old schedule bot with group membership bot

Browse Source
pull/313/head
Mike Lang 2 years ago committed by Mike Lang
parent c286b711ea
commit ebeb9f3f9e
2 changed files with 224 additions and 117 deletions
Show Stats Download Patch File Download Diff File

117
zulip/group_membership.py
Unescape Escape View File

@ -1,117 +0,0 @@
import gevent.monkey
gevent.monkey.patch_all()
import logging
import requests
logging.basicConfig(level='INFO')
class Client(object):
def __init__(self, base_url, email, api_key):
self.base_url = base_url
self.email = email
self.api_key = api_key
self.session = requests.Session()
def request(self, method, *path, **params):
if method == 'GET':
args = {"params": params}
else:
args = {"data": params}
url = "/".join([self.base_url, "api/v1"] + list(path))
resp = session.request(method, url, auth=(self.email, self.api_key), **args)
resp.raise_for_status()
return resp.json()
def get_membership(client):
"""Returns {group id: member set}"""
return {
group["id"]: set(group["members"])
for group in client.request("GET", "user_groups")["user_groups"]
}
def update_members(client, group_id, old_members, new_members):
added = new_members - old_members
removed = old_members - new_members
client.request("POST", "user_groups", group_id, "members", add=list(added), delete=list(removed))
def determine_members(schedules, hour):
return set(
user for user, schedule in schedules.items()
if hour in schedule
)
def run_hour(client, user_map, groups, hour):
logging.info("Setting groups for hour {}".format(hour))
members = get_membership(client)
def run_group(group_id, schedules):
new_members = determine_members(schedules, hour)
new_members = set(user_map[id] for id in new_members)
assert group_id in members, "group {} doesn't exist".format(group_id)
update_members(client, group_id, members[group_id], new_members)
gevent.pool.Group().map(run_group, groups.items())
def parse_config(conf_file):
MAX_DAYS=8
with open(conf_file) as f:
config = yaml.safe_load(f)
assert "url" in config
assert "api_key" in config
assert "start_time" in config
all_users = set()
for group_id, schedules in config["groups"]:
all_users |= set(schedules.keys())
for user_id, schedule in schedules:
hours = set()
for part in scuedule.split(","):
part = part.strip()
if part.endswith("/24"):
hour = int(part.split('/')[0])
hours |= set(24 * day + hour for day in range(MAX_DAYS))
elif part in "ZDAN":
shift = "ZDAN".index(part)
hours |= set(hour for hour in range(24 * MAX_DAYS) if (hour % 24) // 6 == shift)
else:
hours.add(int(part))
schedules[user_id] = hours
missing = all_users - set(config["members"])
assert not missing, "missing: {}".format(", ".join(missing))
def main(conf_file, hour=-1):
"""
config:
url: the base url of the instance
api_key: authentication
start_time: Time of the first hour, as epoch int
members:
NAME: USER_ID
groups:
GROUP_ID:
NAME: SCHEDULE
Where:
schedule: Comma-seperated list of hours
hour: One of:
Integer hour, representing that hour of the run
N/24, expanding to hour N of each day (eg. 0/24 is midnight each day)
Z | D | A | N, expanding to all hours of Zeta, Dawn Guard, Alpha Flight, Night Watch respectively.
"""
config = parse_config(conf_file)
client = Client(config["url"], config["email"], config["api_key"])
user_map = config["members"]
groups = config["groups"]
if hour >= 0:
run_hour(client, user_map, groups, hour)
return
while True:
hour = (time.time() - start_time) // 3600
run_hour(client, user_map, groups, hour)
next_hour = start_time + 3600 * (hour + 1)
remaining = next_hour - time.time()
if remaining > 0:
time.sleep(remaining)
if __name__ == '__main__':
argh.dispatch_command(main)

224
zulip/schedulebot.py
Unescape Escape View File

@ -0,0 +1,224 @@
import gevent.monkey
gevent.monkey.patch_all()
import csv
import json
import logging
import time
from datetime import datetime, timedelta
import gevent.pool
import argh
import requests
session = requests.Session()
logging.basicConfig(level='INFO')
class Client(object):
def __init__(self, base_url, email, api_key):
self.base_url = base_url
self.email = email
self.api_key = api_key
def request(self, method, *path, **params):
if method == 'GET':
args = {"params": params}
else:
args = {"data": {
k: v if isinstance(v, str) else json.dumps(v)
for k, v in params.items()
}}
url = "/".join([self.base_url, "api/v1"] + list(map(str, path)))
resp = session.request(method, url, auth=(self.email, self.api_key), **args)
if not resp.ok:
logging.info(repr(params))
logging.info(f"Got {resp.status_code} for {url}: {resp.text}")
resp.raise_for_status()
return resp.json()
def get_membership(client):
"""Returns {group id: member set}"""
return {
group["id"]: set(group["members"])
for group in client.request("GET", "user_groups")["user_groups"]
}
def update_members(client, group_id, old_members, new_members):
logging.info(f"Updating group {group_id}: {old_members} -> {new_members}")
added = new_members - old_members
removed = old_members - new_members
if added or removed:
client.request("POST", "user_groups", group_id, "members", add=list(added), delete=list(removed))
def get_role_at_hour(hours, hour):
if 0 <= hour < len(hours):
return hours[hour]
return ""
def determine_members(schedule, role, hour):
return set(
user_id for user_id, (_, hours) in schedule.items()
if get_role_at_hour(hours, hour) == role
)
def get_display_name(client, user_id):
return client.request("GET", "users", user_id)["user"]["full_name"]
def update_groups(client, group_ids, schedule, hour):
logging.info("Setting groups for hour {}".format(hour))
members = get_membership(client)
def run_group(item):
group_name, group_id = item
new_members = determine_members(schedule, group_name, hour)
assert group_id in members, "group {} doesn't exist".format(group_id)
update_members(client, group_id, members[group_id], new_members)
gevent.pool.Group().map(run_group, group_ids.items())
def post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions):
going_offline = []
coming_online = []
display_names = {}
supervisor = None
for user_id, (_, hours) in schedule.items():
prev = get_role_at_hour(hours, hour - 1)
now = get_role_at_hour(hours, hour)
if now == "Supervisor":
supervisor = user_id
if user_id not in display_names:
display_names[user_id] = gevent.spawn(get_display_name, client, user_id)
if prev != now:
if prev != "":
going_offline.append((prev, user_id))
if now != "":
coming_online.append((now, user_id))
if user_id not in display_names:
display_names[user_id] = gevent.spawn(get_display_name, client, user_id)
# sort by role
going_offline.sort()
coming_online.sort()
logging.info(f"Going offline: {going_offline}")
logging.info(f"Coming online: {coming_online}")
start_time_pst = datetime.utcfromtimestamp(start_time) - timedelta(hours=8)
current_time = (start_time_pst + timedelta(hours=hour)).replace(minute=0, second=0, microsecond=0)
hour_pst = current_time.hour
shift = hour_pst // 6
shift = ["night-watch", "zeta", "dawn-guard", "alpha-flight"][shift]
shift_hour = hour_pst % 6 + 1
def render_name(user_id, mention=True):
if no_mentions:
mention = False
fallback, _ = schedule[user_id]
try:
result = display_names[user_id].get()
except Exception:
logging.warning(f"Failed to fetch user {user_id}", exc_info=True)
return f"**{fallback}**"
if mention:
return f"@**{result}|{user_id}**"
else:
return f"**{result}**"
lines = [
f"**Shift changes for :{shift}: Hour {shift_hour} | Bustime {hour:02d}:00 - {hour+1:02d}:00 | <time:{current_time.isoformat(timespec='minutes')}>:**",
"Make sure to *mute/unmute* #**current-shift** as needed!",
]
if supervisor is None:
logging.warning("No supervisor found")
else:
name = render_name(supervisor, mention=False)
lines.append(f"Your shift supervisor is {name}")
if coming_online:
lines += [
"",
"---",
"Coming online:"
] + [
f"- {render_name(user_id)} - {role}"
for role, user_id in coming_online
]
if going_offline:
lines += [
"",
"---",
"Going offline:"
] + [
f"- {render_name(user_id)} - {role}"
for role, user_id in going_offline
]
lines += [
"",
"---",
]
send_client.request("POST", "messages",
type="stream",
to=stream,
topic="Schedule",
content="\n".join(lines),
)
def parse_schedule(user_ids, schedule_file):
schedule = {}
with open(schedule_file) as f:
for row in csv.reader(f):
name = row[0]
if name in ["", "Chat Member", "[Counts]", "Hour of the Run"] or name.startswith("-"):
continue
if name not in user_ids:
logging.warning(f"No user id known for user {name}")
continue
user_id = user_ids[name]
schedule[user_id] = name, row[1:]
return schedule
def main(conf_file, hour=-1, no_groups=False, stream="General", no_mentions=False):
"""
config:
url: the base url of the instance
api_user: auth used for general api calls
send_user:
auth used for sending messages
defaults to api_user, but you may want a vanity name / avatar
start_time: Time of the first hour, as epoch int
schedule: Path to the schedule CSV file
members:
NAME: USER_ID
groups:
NAME: GROUP_ID
authentication is an object {email, api_key}
"""
with open(conf_file) as f:
config = json.load(f)
client = Client(config["url"], config["api_user"]["email"], config["api_user"]["api_key"])
send_auth = config.get("send_user", config["api_user"])
send_client = Client(config["url"], send_auth["email"], send_auth["api_key"])
group_ids = config["groups"]
start_time = config["start_time"]
schedule = parse_schedule(config["members"], config["schedule"])
if hour >= 0:
if not no_groups:
update_groups(client, group_ids, schedule, hour)
if stream:
post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions)
return
while True:
hour = (time.time() - start_time) // 3600
if not no_groups:
update_groups(client, group_ids, schedule, hour)
if stream:
post_schedule(client, send_client, start_time, schedule, stream, hour, no_mentions)
next_hour = start_time + 3600 * (hour + 1)
remaining = next_hour - time.time()
if remaining > 0:
time.sleep(remaining)
if __name__ == '__main__':
argh.dispatch_command(main)
Write Preview
Loading…
Cancel
Save