diff --git a/core.py b/core.py index a530524..9ed44eb 100644 --- a/core.py +++ b/core.py @@ -2,114 +2,127 @@ import os import sys import configparser import telegram -import threading import time import strings import worker -# Check if a configuration file exists, create one if it doesn't and get the template version number. -with open("config/template_config.ini") as template_file: - # Check if the config file exists - if not os.path.isfile("config/config.ini"): - # Copy the template file to the config file +def main(): + """The core code of the program. Should be run only in the main process!""" + + # Check if a configuration file exists, create one if it doesn't and get the template version number. + with open("config/template_config.ini") as template_file: + # Check if the config file exists + if not os.path.isfile("config/config.ini"): + # Copy the template file to the config file + with open("config/config.ini", "w") as config_file: + config_file.write(template_file.read()) + # Find the template version number + config = configparser.ConfigParser() + config.read_file(template_file) + template_version = int(config["Config"]["version"]) + + # Overwrite the template config with the values in the config + with open("config/config.ini") as config_file: + config.read_file(config_file) + + # Check if the file has been edited + if config["Config"]["is_template"] == "yes": + print("A config file has been created in config/config.ini.\n" + "Edit it with your configuration, set the is_template flag to false and restart this script.") + sys.exit(1) + + # Check if the version has changed from the template + if template_version > int(config["Config"]["version"]): + # Reset the is_template flag + config["Config"]["is_template"] = "yes" + # Update the config version + config["Config"]["version"] = str(template_version) + # Save the file with open("config/config.ini", "w") as config_file: - config_file.write(template_file.read()) - # Find the template version number - config = configparser.ConfigParser() - config.read_file(template_file) - template_version = int(config["Config"]["version"]) + config.write(config_file) + # Notify the user and quit + print("The config file in config/config.ini has been updated.\n" + "Edit it with the new required data, set the is_template flag to true and restart this script.") + sys.exit(1) -# Overwrite the template config with the values in the config -with open("config/config.ini") as config_file: - config.read_file(config_file) + # Create a bot instance + bot = telegram.Bot(config["Telegram"]["token"]) -# Check if the file has been edited -if config["Config"]["is_template"] == "yes": - print("A config file has been created in config/config.ini.\n" - "Edit it with your configuration, set the is_template flag to false and restart this script.") - sys.exit(1) + # Test the specified token + try: + bot.get_me() + except telegram.error.Unauthorized: + print("The token you have entered in the config file is invalid.\n" + "Fix it, then restart this script.") + sys.exit(1) -# Check if the version has changed from the template -if template_version > int(config["Config"]["version"]): - # Reset the is_template flag - config["Config"]["is_template"] = "yes" - # Update the config version - config["Config"]["version"] = str(template_version) - # Save the file - with open("config/config.ini", "w") as config_file: - config.write(config_file) - # Notify the user and quit - print("The config file in config/config.ini has been updated.\n" - "Edit it with the new required data, set the is_template flag to true and restart this script.") - sys.exit(1) + # Create a dictionary linking the chat ids to the ChatWorker objects + # {"1234": } + chat_workers = {} -# Create a bot instance -bot = telegram.Bot(config["Telegram"]["token"]) + # Current update offset; if None it will get the last 100 unparsed messages + next_update = None -# Test the specified token -try: - bot.get_me() -except telegram.error.Unauthorized: - print("The token you have entered in the config file is invalid.\n" - "Fix it, then restart this script.") - sys.exit(1) - -# Create a dictionary linking the chat ids to the ChatWorker objects -# {"1234": } -chat_workers = {} - -# Current update offset; if None it will get the last 100 unparsed messages -next_update = None + # Main loop of the program + while True: + # Get a new batch of 100 updates and mark the last 100 parsed as read + # TODO: handle possible errors + updates = bot.get_updates(offset=next_update) + # Parse all the updates + for update in updates: + # If the update is a message... + if update.message is not None: + # Ensure the message has been sent in a private chat + if update.message.chat.type != "private": + # Notify the chat + bot.send_message(update.message.chat.id, strings.error_nonprivate_chat) + # Skip the update + continue + # If the message is a start command... + if update.message.text == "/start": + # Check if a worker already exists for that chat + old_worker = chat_workers.get(update.message.chat.id) + # If it exists, gracefully stop the worker + if old_worker: + old_worker.stop() + # Initialize a new worker for the chat + new_worker = worker.ChatWorker(bot, update.message.chat) + # Start the worker + new_worker.start() + # Store the worker in the dictionary + chat_workers[update.message.chat.id] = new_worker + # Skip the update + continue + # Otherwise, forward the update to the corresponding worker + receiving_worker = chat_workers.get(update.message.chat.id) + # Ensure a worker exists for the chat + if receiving_worker is None: + # Suggest that the user restarts the chat with /start + bot.send_message(update.message.chat.id, strings.error_no_worker_for_chat) + # Skip the update + continue + # Forward the update to the worker + receiving_worker.pipe.send(update) + # If the update is a inline keyboard press... + if update.inline_query is not None: + # Forward the update to the corresponding worker + receiving_worker = chat_workers.get(update.inline_query.chat.id) + # Ensure a worker exists for the chat + if receiving_worker is None: + # Suggest that the user restarts the chat with /start + bot.send_message(update.inline_query.chat.id, strings.error_no_worker_for_chat) + # Skip the update + continue + # Forward the update to the worker + receiving_worker.pipe.send(update) + # If there were any updates... + if len(updates): + # Mark them as read by increasing the update_offset + next_update = updates[-1].update_id + 1 + # Temporarily prevent rate limits (remove this later) + time.sleep(5) -# Main loop of the program -while True: - # Get a new batch of 100 updates and mark the last 100 parsed as read - updates = bot.get_updates(offset=next_update) - # Parse all the updates - for update in updates: - # If the update is a message... - if update.message is not None: - # Ensure the message has been sent in a private chat - if update.message.chat.type != "private": - # Notify the chat - bot.send_message(update.message.chat.id, strings.error_nonprivate_chat) - # Skip the update - continue - # If the message is a start command... - if update.message.text == "/start": - # Check if a worker already exists for that chat - old_worker = chat_workers.get(update.message.chat.id) - # If it exists, gracefully stop the worker - if old_worker: - old_worker.stop() - # Initialize a new worker for the chat - new_worker = worker.ChatWorker(bot, update.message.chat) - # Start the worker - new_worker.start() - # Store the worker in the dictionary - chat_workers[update.message.chat.id] = new_worker - # Skip the update - continue - # Otherwise, forward the update to the corresponding worker - receiving_worker = chat_workers.get(update.message.chat.id) - # Ensure a worker exists for the chat - if receiving_worker is None: - # Suggest that the user restarts the chat with /start - bot.send_message(update.message.chat.id, strings.error_no_worker_for_chat) - # Forward the update to the worker - receiving_worker.pipe.send(update) - # If the update is a inline keyboard press... - if update.inline_query is not None: - # Forward the update to the corresponding worker - receiving_worker = chat_workers.get(update.message.chat.id) - # Ensure a worker exists for the chat - if receiving_worker is None: - # Suggest that the user restarts the chat with /start - bot.send_message(update.message.chat.id, strings.error_no_worker_for_chat) - # Forward the update to the worker - receiving_worker.pipe.send(update) - # Mark the update as read by increasing the update_offset - next_update = update.update_id + 1 - # Temporarily prevent rate limits (remove this later) - time.sleep(5) \ No newline at end of file +# Run the main function only in the main process +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/strings.py b/strings.py index c1b65fb..4115cb6 100644 --- a/strings.py +++ b/strings.py @@ -1,9 +1,14 @@ # Strings / localization file for greed # Can be edited, but DON'T REMOVE THE REPLACEMENT FIELDS (words surrounded by {curly braces}) +# TODO: maybe add a preformat to all strings in this file + +# Answer: the start command was sent and the bot should welcome the user +conversation_after_start = "Ciao!\n" \ + "Benvenuto su greed!" # Error: message received not in a private chat error_nonprivate_chat = "⚠️ Questo bot funziona solo in chat private." # Error: a message was sent in a chat, but no worker exists for that chat. Suggest the creation of a new worker with /start error_no_worker_for_chat = "⚠️ La conversazione con il bot si è interrotta.\n" \ - "Per riavviarla, manda il comando /start al bot." \ No newline at end of file + "Per riavviarla, manda il comando /start al bot." diff --git a/worker.py b/worker.py index 34c645d..3a86163 100644 --- a/worker.py +++ b/worker.py @@ -1,12 +1,20 @@ import multiprocessing import telegram +import strings + +class StopSignal: + """A data class that should be sent to the worker when the conversation has to be stopped abnormally.""" + + def __init__(self, reason: str=""): + self.reason = reason + class ChatWorker: """A worker for a single conversation. A new one should be created every time the /start command is sent.""" def __init__(self, bot: telegram.Bot, chat: telegram.Chat): # A pipe connecting the main process to the chat process is created - in_pipe, out_pipe = multiprocessing.Pipe(duplex=False) + out_pipe, in_pipe = multiprocessing.Pipe(duplex=False) # The sending pipe is stored in the ChatWorker class, allowing the forwarding of messages to the chat process self.pipe = in_pipe # A new process running the conversation handler is created, and the receiving pipe is passed to its arguments to enable the receiving of messages @@ -16,13 +24,40 @@ class ChatWorker: """Start the worker process.""" self.process.start() - def stop(self): - # Gracefully stop the worker process - # TODO: send a stop message to the process - raise NotImplementedError() + def stop(self, reason: str=""): + """Gracefully stop the worker process""" + # Send a stop message to the process + self.pipe.send(StopSignal(reason)) # Wait for the process to stop self.process.join() -def conversation_handler(bot: telegram.Bot, chat: telegram.Chat, pipe: multiprocessing.Connection): - raise NotImplementedError() \ No newline at end of file +def receive_next_update(pipe) -> telegram.Update: + """Get the next update from a pipe. + If no update is found, block the process until one is received. + If a stop signal is sent, try to gracefully stop the process.""" + # Receive data from the pipe + data = pipe.recv() + # Check if the data is a stop signal instance + if isinstance(data, StopSignal): + # Gracefully stop the process + graceful_stop() + # Return the received update + return data + + +def graceful_stop(): + """Handle the graceful stop of the process.""" + raise NotImplementedError() + + +def conversation_handler(bot: telegram.Bot, chat: telegram.Chat, pipe): + """This function is ran once for every conversation (/start command) by a separate process.""" + # TODO: catch all the possible exceptions + # Welcome the user to the bot + bot.send_message(chat.id, strings.conversation_after_start) + # TODO: Send a command list or something + while True: + # For now, echo the sent message + update = receive_next_update(pipe) + bot.send_message(chat.id, update.message.text) \ No newline at end of file