mostly State updates

2 files changed, 30 insertions, 14 deletions

diff --git a/ch1p/functions.py b/ch1p/functions.py
index cf75221..7cc2a10 100644
--- a/ch1p/functions.py
+++ b/ch1p/functions.py
@@ -13,7 +13,7 @@ def _get_vars(params: List[Tuple], kw: dict) -> List[AnyStr]:
         elif env_name in os.environ:
             result.append(os.environ[env_name])
         else:
-            raise RuntimeError("missing parameter %s or variable %s" % (kw_name, env_name))
+            raise RuntimeError("missing parameter '%s' or variable '%s'" % (kw_name, env_name))
 
     return result
 
diff --git a/ch1p/state.py b/ch1p/state.py
index e7fe050..1ea2373 100644
--- a/ch1p/state.py
+++ b/ch1p/state.py
@@ -1,28 +1,44 @@
-import os, json
+import os, json, atexit
 
 from .functions import _get_vars
 
 
 class State:
-    def __init__(self, default=None, **kwargs):
+    def __init__(self, default: dict = None, **kwargs):
         file, = _get_vars([
             ('file', 'STATE_FILE')
         ], kwargs)
 
         if default is None:
             default = {}
+        elif type(default) is not dict:
+            raise TypeError('default must be dictionary')
 
-        self.file = file
-        self.default = default
+        if not os.path.exists(file):
+            self._data = default
+        else:
+            with open(file, 'r') as f:
+                self._data = json.loads(f.read())
 
-    def read(self) -> dict:
-        if not os.path.exists(self.file):
-            self.write(self.default)
-            return self.default
+        self._file = file
+        atexit.register(self.__cleanup)
 
-        with open(self.file, 'r') as f:
-            return json.loads(f.read())
+    def __cleanup(self):
+        if hasattr(self, '_file'):
+            with open(self._file, 'w') as f:
+                f.write(json.dumps(self._data))
+        atexit.unregister(self.__cleanup)
 
-    def write(self, state: dict):
-        with open(self.file, 'w') as f:
-            f.write(json.dumps(state))
\ No newline at end of file
+    def __del__(self):
+        if 'open' in __builtins__:
+            self.__cleanup()
+
+    def __getitem__(self, key):
+        return self._data[key]
+
+    def __setitem__(self, key, value):
+        self._data[key] = value
+
+    def __delitem__(self, key):
+        if key in self._data:
+            del self._data[key]
\ No newline at end of file