diff --git a/.gitignore b/.gitignore index 51272fe..1d76580 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ venv __pycache__ *.session +data/ diff --git a/uniborg/storage.py b/uniborg/storage.py new file mode 100644 index 0000000..d19f773 --- /dev/null +++ b/uniborg/storage.py @@ -0,0 +1,43 @@ +import os +import json + + +class Storage: + class _Guard: + def __init__(self, storage): + self._storage = storage + + def __enter__(self): + self._storage._autosave = False + + def __exit__(self, *args): + self._storage._autosave = True + self._storage._save() + + def __init__(self, root): + self._root = root + self._autosave = True + self._guard = self._Guard(self) + self._data = {} + + def bulk_save(self): + return self._guard + + def __getattr__(self, name): + if name.startswith('_'): + raise ValueError('You can only access existing private members') + return self._data.get(name, None) + + def __setattr__(self, name, value): + if name.startswith('_'): + self.__dict__[name] = value + else: + self._data[name] = value + if self._autosave: + self._save() + + def _save(self): + if not os.path.isdir(self._root): + os.makedirs(self._root, exist_ok=True) + with open(os.path.join(self._root, 'data.json'), 'w') as fp: + json.dump(self._data, fp) diff --git a/uniborg/telethon.py b/uniborg/telethon.py index 812457d..0d3a2d7 100644 --- a/uniborg/telethon.py +++ b/uniborg/telethon.py @@ -1,7 +1,7 @@ # This Source Code Form is subject to the terms of the Mozilla Public # License, v. 2.0. If a copy of the MPL was not distributed with this # file, You can obtain one at http://mozilla.org/MPL/2.0/. - +import os import asyncio import importlib.util import logging @@ -11,15 +11,20 @@ from telethon import TelegramClient import telethon.utils import telethon.events +from .storage import Storage from . import hacks class Uniborg(TelegramClient): def __init__( - self, session, *, plugin_path="plugins", + self, session, *, plugin_path="plugins", storage=None, bot_token=None, **kwargs): # TODO: handle non-string session + # + # storage should be a callable accepting plugin name -> Storage object. + # This means that using the Storage type as a storage would work too. self._name = session + self.storage = storage or (lambda n: Storage(os.path.join('data', n))) self._logger = logging.getLogger(session) self._plugins = {} self._plugin_path = plugin_path @@ -64,6 +69,7 @@ class Uniborg(TelegramClient): mod.borg = self mod.logger = logging.getLogger(shortname) + mod.storage = self.storage(shortname) spec.loader.exec_module(mod) self._plugins[shortname] = mod