Added storage capabilities to pgzero

doc/builtins.rst

@@ -738,3 +738,63 @@ This could be used in a Pygame Zero program like this::
 
     def on_mouse_down():
         beep.play()
+
+
+.. _data_storage:
+
+
+Data Storage
+------------
+
+The ``storage`` object behaves just like a dictionary but has two additional methods to
+allow to save/load the data to/from the disk.
+
+Because it's a dictionary-like object, it supports all operations a dictionary does.
+For example, you can update the storage with another dictionary, like so::
+
+    my_data = {'player_turn': 1, 'level': 10}
+    storage.update(my_data)
+
+On windows the data is saved under ``%APPDATA%/pgzero/saves/`` and on Linux/MacOS under ``~/.config/pgzero/saves/``.
+
+The saved files will be named after their module name.
+
+**NOTE:** Make sure your scripts have different names, otherwise they will be picking each other data.
+
+.. class:: Storage
+
+    .. method:: save()
+
+        Saves the data to disk.
+
+    .. method:: load()
+
+        Loads the data from disk.
+
+Example of usage::
+
+    # Setting some values
+    storage['my_score'] = 500
+    storage['level'] = 1
+
+    # You can have nested lists and dictionaries
+    storage['high_scores'] = []
+    storage['high_scores'].append(10)
+    storage['high_scores'].append(12)
+    storage['high_scores'].append(11)
+    storage['high_scores'].sort()
+
+    # Save storage to disk.
+    storage.save()
+
+
+Following on the previous example, when starting your program, you can load that data back in::
+
+    storage.load()
+
+    my_score = storage['my_score']
+
+    level = storage['level']
+
+    # Can use the get method from dicts to return a default value
+    storage.get('lifes', 3)

pgzero/builtins.py

@@ -3,6 +3,7 @@
 from . import music
 from . import tone
 from .actor import Actor
+from .storage import storage
 from .keyboard import keyboard
 from .animation import animate
 from .rect import Rect, ZRect

pgzero/runner.py

@@ -12,6 +12,7 @@
 from .game import PGZeroGame, DISPLAY_FLAGS
 from . import loaders
 from . import builtins
+from .storage import Storage
 
 
 def _check_python_ok_for_pygame():
@@ -102,6 +103,8 @@ def prepare_mod(mod):
     set before the module globals are run.
 
     """
+    fn_hash = hash(mod.__file__) % ((sys.maxsize + 1) * 2)
+    Storage.set_app_hash(format(fn_hash, 'x'))
     loaders.set_root(mod.__file__)
     PGZeroGame.show_default_icon()
     pygame.display.set_mode((100, 100), DISPLAY_FLAGS)

pgzero/storage.py

@@ -0,0 +1,114 @@
+import json
+import os
+import platform
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = ['StorageCorruptionException', 'JSONEncodingException', 'Storage']
+
+
+class StorageCorruptionException(Exception):
+    """The data in the storage is corrupted."""
+
+
+class JSONEncodingException(Exception):
+    """The data in the storage is corrupted."""
+
+
+class Storage(dict):
+
+    path = ''
+
+    """Behaves like a dictionary with a few extra functions.
+
+    It's possible to load/save the data to disk.
+
+    The name of the file will be the script's name hashed.
+
+    NOTE: If two scripts have the same name, they will load/save from/to
+    the same file.
+    """
+    def __init__(self):
+        dict.__init__(self)
+
+        storage_path = self._get_platform_pgzero_path()
+        if not os.path.exists(storage_path):
+            os.makedirs(storage_path)
+
+    @classmethod
+    def set_app_hash(cls, name):
+        storage_path = cls._get_platform_pgzero_path()
+        cls.path = os.path.join(storage_path, '{}.json'.format(name))
+
+    def load(self):
+        """Load data from disk."""
+        self.clear()
+
+        try:
+            with open(self.path, 'r') as f:
+                data = json.load(f)
+        except FileNotFoundError:
+            # If no file exists, it's fine
+            logger.debug('No file to load data from.')
+        except json.JSONDecodeError:
+            msg = 'Storage is corrupted. Couldn\'t load the data.'
+            logger.error(msg)
+            raise StorageCorruptionException(msg)
+        else:
+            self.update(data)
+
+    def save(self):
+        """Save data to disk."""
+        try:
+            data = json.dumps(self)
+        except TypeError:
+            json_path, key_type_obj = self._get_json_error_keys(self)
+            msg = 'The following entry couldn\'t be JSON serialized: storage{} of type {}'
+            msg = msg.format(json_path, key_type_obj)
+            logger.error(msg)
+            raise JSONEncodingException(msg)
+        else:
+            with open(self.path, 'w+') as f:
+                f.write(data)
+
+    @classmethod
+    def _get_json_error_keys(cls, obj, json_path=''):
+        """Work out which part of the storage failed to be JSON encoded.
+
+        This function returns None if there is no error.
+        """
+
+        if type(obj) in (float, int, str, bool):
+            return None
+        elif isinstance(obj, list):
+            for i, value in enumerate(obj):
+                result = cls._get_json_error_keys(value, '{}[{}]'.format(json_path, i))
+                if result is not None:
+                    return result
+        elif isinstance(obj, dict):
+            for k, v in obj.items():
+                result = cls._get_json_error_keys(v, '{}[\'{}\']'.format(json_path, k))
+                if result is not None:
+                    return result
+        else:
+            # TODO: Could have a custom JSONEncoder in the future, in which case we would
+            # have to actually run json.dumps on it to get the result and possibly even
+            # drill down on the parameters we got back, if that object happens to have
+            # a list as an attribute, for example.
+            return (json_path, type(obj))
+
+        # Explicitly returning None for the sake of code readability
+        return None
+
+    @staticmethod
+    def _get_platform_pgzero_path():
+        """Attempts to get the right directory in Windows/Linux.MacOS"""
+        if platform.system() == 'Windows':
+            return os.path.join(os.environ['APPDATA'], 'pgzero')
+        return os.path.expanduser(os.path.join('~', '.config/pgzero/saves'))
+
+
+storage = Storage()

test/test_storage.py

@@ -0,0 +1,42 @@
+import unittest
+from unittest import mock
+from unittest.mock import patch
+
+from pgzero.storage import Storage
+
+
+class StorageStaticMethodsTest(unittest.TestCase):
+    def test_dict_with_no_errors(self):
+        obj = {'level': 10, 'player_name': 'Daniel'}
+
+        result = Storage._get_json_error_keys(obj)
+        self.assertIsNone(result)
+
+    def test_dict_with_errors(self):
+        obj = {'level': 10, 'player_name': 'Daniel', 'obj': object()}
+        expected_result = ("['obj']", type(object()))
+
+        result = Storage._get_json_error_keys(obj)
+        self.assertEqual(result, expected_result)
+
+    def test_dict_with_nested_dicts_errors(self):
+        subobj0 = {'this_key_fails': object()}
+        subobj1 = {'a': 10, 'b': 20, 'c': 30, 'obj': subobj0}
+        subobj2 = {'player_name': 'Daniel', 'level': 20, 'states': subobj1}
+        obj = {'game': 'my_game', 'state': subobj2}
+        expected_result = ("['state']['states']['obj']['this_key_fails']", type(object()))
+
+        result = Storage._get_json_error_keys(obj)
+        self.assertEqual(result, expected_result)
+
+
+class StorageTest(unittest.TestCase):
+    @patch('pgzero.storage.os.path.exists')
+    def setUp(self, exists_mock):
+        exists_mock.return_value = True
+        self.storage = Storage()
+
+    def test_get(self):
+        with patch('builtins.open', mock.mock_open(read_data='{"a": "hello"}')) as m:
+            self.storage.load()
+            self.assertEqual(self.storage['a'], 'hello')