replit-py/src/replit/database/jsonkey.py

# flake8: noqa
from typing import Any, Callable, Dict, Optional, Tuple, Union
import json

JSON_TYPE = Optional[Union[str, int, float, bool, dict, list]]


class AsyncJSONKey:
    """Represents an key in the async database that holds a JSON value.
    db.jsonkey() will initialize an instance for you,
    you don't have to do it manually.
    """

    __slots__ = ("db", "key", "dtype", "get_default", "discard_bad_data", "do_raise")

    def __init__(
        self,
        db: Any,
        key: str,
        dtype: JSON_TYPE,
        get_default: Callable[[], JSON_TYPE] = None,
        discard_bad_data: bool = False,
        do_raise: bool = False,
    ) -> None:
        """Initialize the key.
        Args:
            db (Any): An instance of ReplitDb
            key (str): The key to read
            dtype (JSON_TYPE): The datatype the key should be, can be typing.Any.
            get_default (Callable): A function that returns the default
                value if the key is not set. If it is None (the default) the dtype
                argument is used.
            discard_bad_data (bool): Don't prompt if bad data is read, overwrite it
                with the default. Defaults to False.
            do_raise (bool): Whether to raise exceptions when errors are encountered.
        """
        self.db = db
        self.key = key
        self.dtype = dtype
        self.get_default = get_default
        self.discard_bad_data = discard_bad_data
        self.do_raise = do_raise

    def _default(self) -> JSON_TYPE:
        if self.get_default:
            return self.get_default()
        return self.dtype

    def _is_valid_type(self, data: JSON_TYPE) -> bool:
        return self.dtype is Any or isinstance(data, self.dtype)

    def _type_mismatch_msg(self, data: Any) -> str:
        return (
            f"Type mismatch: Got type {type(data).__name__},"
            f"expected {self.dtype.__name__}"
        )

    async def get(self) -> JSON_TYPE:
        """Get the value of the key.
        If an invalid JSON value is read or the type does not match, it will show a
            prompt asking the user what to do unless discard_bad_data is set.
        Raises:
            KeyError: If do_raise is true and the key does not exist.
            json.JSONDecodeError: If do_raise is true and invalid JSON data is read
                from the key.
        Returns:
            JSON_TYPE: The value read from the database
        """
        try:
            read = await self.db.get(self.key)
        except KeyError:
            if self.do_raise:
                raise
            print(f"Database key {self.key} not set, setting it to default value")
            default = self._default()
            await self.db.set(self.key, default)
            return default

        try:
            data = json.loads(read)
        except json.JSONDecodeError:
            if self.do_raise:
                raise
            return await self._error("Invalid JSON data read", read)

        if not self._is_valid_type(data):
            return await self._error(self._type_mismatch_msg(data), read,)
        return data

    async def _error(self, error: str, read: str) -> JSON_TYPE:
        if self.do_raise:
            raise ValueError(error)

        print(f"Error reading key {self.key!r}: {error}", file=stderr)
        if self.discard_bad_data:
            val = self._default()
            await self.db.set(self.key, json.dumps(val))
            print(f"Wrote default to key {self.key!r}")
            return val
        return await self._should_discard_prompt(error, read)

    async def _should_discard_prompt(self, error: str, read: str) -> bool:
        while True:
            choice = input(
                "d to use default, v to view the invalid data, c to insert custom "
                "value, ^C to exit: "
            )
            if choice.startswith("d"):
                print("Writing default...")
                val = self._default()
                await self.db.set(self.key, val)
                return val
            elif choice.startswith("v"):
                print(f"Data read from key: {read!r}")
            elif choice.startswith("c"):
                toset = input(
                    f"Enter data to write, should be of type {self.dtype.__name__!r}"
                    " (leave blank to return to menu): "
                )
                if not toset:
                    continue
                try:
                    data = json.loads(toset)
                except json.JSONDecodeError:
                    print("Invalid JSON data!")
                else:
                    if not self._is_valid_type(data):
                        print(self._type_mismatch_msg(data))
                        continue

                    await self.db.set(self.key, toset)
                    print("Wrote data to key")
                    return data

    async def set(self, data: JSON_TYPE) -> None:
        """Set the value of the jsonkey.
        Args:
            data (JSON_TYPE): The value to set it to.
        Raises:
            TypeError: The type of the value set does not match the datatype.
        """
        if not self._is_valid_type(data):
            raise TypeError(self._type_mismatch_msg(data))

        await self.db.set(self.key, json.dumps(data))


class JSONKey(AsyncJSONKey):
    """Represents a key in the database that holds a JSON value.

    db.jsonkey() will initialize an instance for you,
    you don't have to do it manually.
    """

    __slots__ = ("db", "key", "dtype", "get_default", "discard_bad_data", "do_raise")

    def __init__(
        self,
        db: Any,
        key: str,
        dtype: JSON_TYPE,
        get_default: Callable = None,
        discard_bad_data: bool = False,
        do_raise: bool = False,
    ) -> None:
        """Initialize the key.

        Args:
            db (Any): An instance of ReplitDb
            key (str): The key to read
            dtype (JSON_TYPE): The datatype the key should be, can be typing.Any.
            get_default (Callable): A function that returns the default
                value if the key is not set. If it is None (the default) the dtype
                argument is used.
            discard_bad_data (bool): Don't prompt if bad data is read, overwrite it
                with the default. Defaults to False.
            do_raise (bool): Whether to raise exceptions when errors are encountered.
        """
        self.db = db
        self.key = key
        self.dtype = dtype
        self.get_default = get_default
        self.discard_bad_data = discard_bad_data
        self.do_raise = do_raise

    def get(self) -> JSON_TYPE:
        """Get the value of the key.

        If an invalid JSON value is read or the type does not match, it will show a
            prompt asking the user what to do unless discard_bad_data is set.

        Returns:
            JSON_TYPE: The value read from the database
        """
        try:
            read = self.db[self.key]
        except KeyError:
            print(f"Database key {self.key} not set, setting it to default value")
            default = self._default()
            self.db[self.key] = default
            return default

        # TODO: ReplitDb isn't defined here, so this will throw an exception.
        if isinstance(self.db, ReplitDb):
            try:
                data = json.loads(read)
            except json.JSONDecodeError:
                return self._error("Invalid JSON data read", read)
        else:
            data = read

        if not self._is_valid_type(data):
            return self._error(self._type_mismatch_msg(data), read,)
        return data

    def _error(self, error: str, read: str) -> JSON_TYPE:
        print(f"Error reading key {self.key!r}: {error}", file=stderr)
        if self.discard_bad_data:
            val = self._default()
            self.db[self.key] = json.dumps(val)
            print(f"Wrote default to key {self.key!r}")
            return val
        return self._should_discard_prompt(error, read)

    def _should_discard_prompt(self, error: str, read: str) -> bool:
        while True:
            choice = input(
                "d to use default, v to view the invalid data, c to insert custom "
                "value, ^C to exit: "
            )
            if choice.startswith("d"):
                print("Writing default...")
                val = self._default()
                self.db[self.key] = val
                return val
            elif choice.startswith("v"):
                print(f"Data read from key: {read!r}")
            elif choice.startswith("c"):
                toset = input(
                    f"Enter data to write, should be of type {self.dtype.__name__!r}"
                    " (leave blank to return to menu): "
                )
                if not toset:
                    continue
                try:
                    data = json.loads(toset)
                except json.JSONDecodeError:
                    print("Invalid JSON data!")
                else:
                    if not self._is_valid_type(data):
                        print(self._type_mismatch_msg(data))
                        continue

                    self.db[self.key] = toset
                    print("Wrote data to key")
                    return data

    def set(self, data: JSON_TYPE) -> None:
        """Set the value of the jsonkey.

        Args:
            data (JSON_TYPE): The value to set it to.

        Raises:
            TypeError: The type of the value set does not match the datatype.
        """
        if not self._is_valid_type(data):
            raise TypeError(self._type_mismatch_msg(data))
        if isinstance(self.db, ReplitDb):
            data = json.dumps(data)
        self.db[self.key] = data

    def read(self, key: str, default: Any = None) -> Any:
        """Shorthand for self.get().get(name, default) if datatype is dict.

        Args:
            key (str): The name to get.
            default (Any): The default if the key doesn't exist. Defaults to None.

        Returns:
            Any: The value read or the default.
        """
        data = self.get()
        if not isinstance(data, dict):
            raise TypeError
        return data.get(key, default)

    def keys(self, *keys: str) -> Any:
        """Reads multiple keys from the key's value and allows setting.

        Args:
            *keys (str): The keys to read from the data.

        Returns:
            Any: The value accessed from self.get()[k1][k2][kn]
        """
        data = self
        for key in keys[:-1]:
            data = type(self)(db=data, key=key, dtype=Any)
        check = data[keys[-1]]
        if type(check) is dict:
            return type(self)(db=data, key=keys[-1], dtype=dict)
        else:
            return check

    def __getitem__(self, name: str) -> JSON_TYPE:
        """Retrieve a key from the JSONKey's value if it is a dict.

        Args:
            name (str): The name to retrieve.

        Returns:
            JSON_TYPE: The value of the key.
        """
        return self.keys(name)

    def __setitem__(self, name: str, value: JSON_TYPE) -> None:
        """Sets a key inside the JSONKey's value if it is a dict.

        Args:
            name (str): The key to set.
            value (JSON_TYPE): The value to set it to.
        """
        data = self.get()
        if not isinstance(data, dict):
            raise TypeError
        data[name] = value
        self.set(data)

    def append(self, item: JSON_TYPE) -> None:
        """Append to the JSONKey's value if it is a list.

        Args:
            item (JSON_TYPE): The item to append.
        """
        data = self.get()
        if not isinstance(data, list):
            raise TypeError
        self.set(data + [item])

    def __add__(self, item: Any) -> Any:
        """Add to the JSONKey's value and return the result.

        Args:
            item (Any): The item to add.

        Returns:
            Any: The result of adding.
        """
        return self.get() + item

    def __iadd__(self, item: Any) -> Any:
        """Add to the JSONKey's value and set the result.

        Args:
            item (Any): The item to add.

        Returns:
            Any: self
        """
        r = self.get() + item
        self.set(r)
        return self