replace pickle with json

This commit is contained in:
Tomáš Mládek 2021-05-26 20:35:37 +02:00
parent 423142af3f
commit 4d29570a2d

28
bot.py
View file

@ -1,7 +1,7 @@
import asyncio import asyncio
import logging import logging
import os import os
import pickle import json
import re import re
import textwrap import textwrap
from asyncio import sleep from asyncio import sleep
@ -17,13 +17,7 @@ config.WEBHOOK_URL = os.getenv("WEBHOOK_URL", config.WEBHOOK_URL)
config.TG_CHAT_ID = os.getenv("TG_CHAT_ID", config.TG_CHAT_ID) config.TG_CHAT_ID = os.getenv("TG_CHAT_ID", config.TG_CHAT_ID)
config.TG_API_KEY = os.getenv("TG_API_KEY", config.TG_API_KEY) config.TG_API_KEY = os.getenv("TG_API_KEY", config.TG_API_KEY)
STATE_PATH = os.getenv("STATE_PATH", "state.pickle") STATE_PATH = os.getenv("STATE_PATH", "state.json")
class State:
def __init__(self) -> None:
self.notified = []
self.notified_soon = []
def process_event(event, template): def process_event(event, template):
@ -100,10 +94,10 @@ def get_event_id(event: Event) -> str:
async def main(): async def main():
try: try:
with open(STATE_PATH, "rb") as state_fp: with open(STATE_PATH, "r") as state_fp:
state = pickle.load(state_fp) state = json.load(state_fp)
except FileNotFoundError: except FileNotFoundError:
state = State() state = {"notified": [], "notified_soon": []}
while True: while True:
calendars = [] calendars = []
@ -123,23 +117,23 @@ async def main():
for event in future_events: for event in future_events:
if ( if (
get_event_id(event) not in state.notified get_event_id(event) not in state["notified"]
and event.begin.shift(days=-2) < arrow.now() and event.begin.shift(days=-2) < arrow.now()
): ):
logging.info(f"Sending description of {event.name}") logging.info(f"Sending description of {event.name}")
await send_message(process_event(event, config.EVENT_TEMPLATE)) await send_message(process_event(event, config.EVENT_TEMPLATE))
state.notified.append(get_event_id(event)) state["notified"].append(get_event_id(event))
if ( if (
get_event_id(event) not in state.notified_soon get_event_id(event) not in state["notified_soon"]
and event.begin.shift(hours=-2) < arrow.now() and event.begin.shift(hours=-2) < arrow.now()
): ):
logging.info(f"Notifying of {event.name}") logging.info(f"Notifying of {event.name}")
await send_message(process_event(event, config.SOON_TEMPLATE)) await send_message(process_event(event, config.SOON_TEMPLATE))
state.notified_soon.append(get_event_id(event)) state["notified_soon"].append(get_event_id(event))
with open(STATE_PATH, "wb") as state_fp: with open(STATE_PATH, "w") as state_fp:
pickle.dump(state, state_fp) json.dump(state, state_fp)
logging.debug("Sleeping for 60s...") logging.debug("Sleeping for 60s...")
await sleep(60) await sleep(60)