From 15169e0bb681f4a4d37f249805ee6342336d132d Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 23 Apr 2018 13:19:55 +0200
Subject: [PATCH] Create a Storage class and inject storage on plugins

---
 .gitignore          |  1 +
 uniborg/storage.py  | 43 +++++++++++++++++++++++++++++++++++++++++++
 uniborg/telethon.py | 10 ++++++++--
 3 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 uniborg/storage.py

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