diff --git a/classes/importer/abstract.py b/classes/importer/abstract.py new file mode 100644 index 0000000..096ec11 --- /dev/null +++ b/classes/importer/abstract.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class Importer(ABC): + """ + The Importer class represents the object with + functionality to import/export garbage collection + records and convert them to other object types. + """ + + @abstractmethod + async def import_data(self, data: Any) -> None: + pass + + @abstractmethod + async def export_data(self, data: Any) -> None: + pass diff --git a/classes/importer/csv.py b/classes/importer/csv.py new file mode 100644 index 0000000..c876f92 --- /dev/null +++ b/classes/importer/csv.py @@ -0,0 +1,64 @@ +from codecs import decode +from csv import DictReader +from datetime import datetime +from typing import Any, Dict, List, Union + +from bson import ObjectId + +from classes.importer.abstract import Importer +from modules.database import col_entries + + +class ImporterCSV(Importer): + """ + The ImporterCSV class represents the object with + functionality to import/export garbage collection + records and convert them to other object types + from CSV files. + """ + + def __init__(self): + super(Importer, self).__init__() + + async def import_data(self, data: bytes) -> List[ObjectId]: + entries: List[Dict[str, Any]] = list( + DictReader(decode(data).split("\n"), delimiter=";") + ) + + for entry in entries: + entry["locations"] = ( + [int(entry["locations"])] + if "," not in entry["locations"] + else [int(id) for id in entry["locations"].split(",")] + ) + entry["garbage_type"] = int(entry["garbage_type"]) + + for key in ("locations", "garbage_type", "date"): + if ( + key not in entry + or (key == "garbage_type" and not isinstance(entry[key], int)) + or (key == "locations" and not isinstance(entry[key], list)) + ): + raise ValueError + + if key == "date": + try: + datetime.fromisoformat(str(entry[key])) + except (ValueError, TypeError) as exc: + raise ValueError from exc + + entries_clean: List[Dict[str, Union[str, int, datetime]]] = [ + { + "locations": entry["locations"], + "garbage_type": entry["garbage_type"], + "date": datetime.fromisoformat(str(entry["date"])), + } + for entry in entries + ] + + inserted = await col_entries.insert_many(entries_clean) + + return [] if inserted is None else inserted.inserted_ids + + async def export_data(self, data: Any) -> Any: + return None diff --git a/classes/importer/json.py b/classes/importer/json.py new file mode 100644 index 0000000..86afaad --- /dev/null +++ b/classes/importer/json.py @@ -0,0 +1,56 @@ +from datetime import datetime +from typing import Any, Dict, List, Union + +from bson import ObjectId +from ujson import loads + +from classes.importer.abstract import Importer +from modules.database import col_entries + + +class ImporterJSON(Importer): + """ + The ImporterJSON class represents the object with + functionality to import/export garbage collection + records and convert them to other object types + from JSON files. + """ + + def __init__(self): + super(Importer, self).__init__() + + async def import_data(self, data: bytes) -> List[ObjectId]: + entries: List[Dict[str, Any]] = loads(data) + + for entry in entries: + for key in ("locations", "garbage_type", "date"): + if ( + key not in entry + or (key == "garbage_type" and not isinstance(entry[key], int)) + or (key == "locations" and not isinstance(entry[key], list)) + ): + print("keys", entry) + raise ValueError + + if key == "date": + try: + datetime.fromisoformat(str(entry[key])) + except (ValueError, TypeError) as exc: + print("date", entry) + raise ValueError from exc + + entries_clean: List[Dict[str, Union[str, int, datetime]]] = [ + { + "locations": entry["locations"], + "garbage_type": entry["garbage_type"], + "date": datetime.fromisoformat(str(entry["date"])), + } + for entry in entries + ] + + inserted = await col_entries.insert_many(entries_clean) + + return [] if inserted is None else inserted.inserted_ids + + async def export_data(self, data: Any) -> Any: + return None diff --git a/locale/en.json b/locale/en.json index 1509a69..b555e9e 100644 --- a/locale/en.json +++ b/locale/en.json @@ -42,9 +42,9 @@ "help": "šŸ”” This bot sends you notifications about garbage collection according to your local schedule.\n\n**Available commands**\n/help - Show this message\n/setup - Select the location\n/toggle - Disable/enable the reminders\n/set_time - Set the reminders' time\n/set_offset - Set offset between reminders and collection\n/upcoming - Show the upcoming collection\n/language - Select the bot's language\n/checkout - Export or remove your data\n\nšŸ’­ You can also suggest adding your town/district to the bot by contacting the admins using [this link]({url_contact}) and providing your schedule.\n\nāš™ļø Want to host this bot yourself or make some changes? It's open-source, so you can basically fork it. Take a look at [bot's repository]({url_repo}) for details.\n\nHappy using! šŸ¤—", "import_finished": "You have successfully inserted {count} entries.", "import_invalid_date": "Entries contain invalid date formats. Use **ISO 8601** date format.", - "import_invalid_filetype": "Invalid input. Please, send me a JSON file with entries. {cancel_notice}", - "import_invalid": "This is not a valid garbage collection JSON.", - "import": "Alright. Send me a valid JSON.", + "import_invalid_filetype": "Invalid input. Please, send me a JSON or CSV file with entries. {cancel_notice}", + "import_invalid": "This is not a valid garbage collection file.", + "import": "Alright. Send me a valid file. It can be in JSON or CSV format. Read more about supported formats in the documentation", "locale_choice": "Alright. Please choose the language using keyboard below.", "location_empty": "You have no location set. Use /setup to select your location first.", "location_name_empty": "Could not find any locations by this name. Try rephrasing it or make sure you use the same location language and name itself as it in written by your local authorities in garbage collection schedule.\n\n{cancel_notice}", diff --git a/plugins/commands/import.py b/plugins/commands/import.py index 64c38e1..9a68c56 100644 --- a/plugins/commands/import.py +++ b/plugins/commands/import.py @@ -1,33 +1,26 @@ -from datetime import datetime -from typing import List, Mapping, Union - from convopyro import listen_message from pyrogram import filters from pyrogram.types import ForceReply, Message, ReplyKeyboardRemove -from ujson import loads +from classes.importer.csv import ImporterCSV +from classes.importer.json import ImporterJSON from classes.pyroclient import PyroClient from modules import custom_filters -from modules.database import col_entries @PyroClient.on_message( - ~filters.scheduled & filters.private & custom_filters.owner & filters.command(["import"], prefixes=["/"]) & ~custom_filters.context # type: ignore + ~filters.scheduled & filters.private & custom_filters.owner & filters.command(["import"], prefixes=["/"]) # type: ignore ) async def command_import(app: PyroClient, message: Message): user = await app.find_user(message.from_user) await message.reply_text( app._("import", "messages", locale=user.locale), - reply_markup=ForceReply( - placeholder=app._("import", "force_replies", locale=user.locale) - ), + reply_markup=ForceReply(placeholder=""), ) while True: - app.contexts.append(message.from_user.id) answer = await listen_message(app, message.chat.id, 300) - app.contexts.remove(message.from_user.id) if answer is None or answer.text == "/cancel": await message.reply_text( @@ -36,7 +29,10 @@ async def command_import(app: PyroClient, message: Message): ) return - if answer.document is None or answer.document.mime_type != "application/json": + if answer.document is None or answer.document.mime_type not in [ + "application/json", + "text/csv", + ]: await answer.reply_text( app._("import_invalid_filetype", "messages", locale=user.locale).format( cancel_notice=app._("cancel", "messages", locale=user.locale) @@ -48,51 +44,36 @@ async def command_import(app: PyroClient, message: Message): file = await app.download_media(answer, in_memory=True) - entries: List[Mapping[str, Union[str, int]]] = loads(bytes(file.getbuffer())) # type: ignore + data: bytes = bytes(file.getbuffer()) # type: ignore - for entry in entries: - if not isinstance(entries, list): - await answer.reply_text( - app._("import_invalid", "messages", locale=user.locale), - reply_markup=ReplyKeyboardRemove(), - ) - return + # I'd like to replace it with switch-case, but 3.9 compatibility + # is still more important to be there. Although refactor may be + # done in the near future as Python 3.9 EOL gets nearer. + if answer.document.mime_type == "application/json": + importer = ImporterJSON() + elif answer.document.mime_type == "text/csv": + importer = ImporterCSV() + else: + await answer.reply_text( + app._("import_invalid_filetype", "messages", locale=user.locale).format( + cancel_notice="" + ), + reply_markup=ReplyKeyboardRemove(), + ) + return - for key in ("locations", "garbage_type", "date"): - if ( - key not in entry - or (key == "garbage_type" and not isinstance(entry[key], int)) - or (key == "locations" and not isinstance(entry[key], list)) - ): - await answer.reply_text( - app._("import_invalid", "messages", locale=user.locale), - reply_markup=ReplyKeyboardRemove(), - ) - return - if key == "date": - try: - datetime.fromisoformat(str(entry[key])) - except (ValueError, TypeError): - await answer.reply_text( - app._("import_invalid_date", "messages", locale=user.locale), - reply_markup=ReplyKeyboardRemove(), - ) - return - - entries_clean: List[Mapping[str, Union[str, int, datetime]]] = [ - { - "locations": entry["locations"], - "garbage_type": entry["garbage_type"], - "date": datetime.fromisoformat(str(entry["date"])), - } - for entry in entries - ] - - await col_entries.insert_many(entries_clean) + try: + import_result = await importer.import_data(data) + except ValueError: + await answer.reply_text( + app._("import_invalid", "messages", locale=user.locale), + reply_markup=ReplyKeyboardRemove(), + ) + return await answer.reply_text( app._("import_finished", "messages", locale=user.locale).format( - count=len(entries_clean) + count=len(import_result) ), reply_markup=ReplyKeyboardRemove(), )