Create a Storage class and inject storage on plugins

ninja_fix
Lonami Exo 2018-04-23 13:19:55 +02:00
parent cf36b5bd4a
commit 15169e0bb6
3 changed files with 52 additions and 2 deletions

1
.gitignore vendored
View File

@ -2,3 +2,4 @@ venv
__pycache__
*.session
data/

43
uniborg/storage.py Normal file
View File

@ -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)

View File

@ -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