summaryrefslogtreecommitdiff
path: root/redis/commands/json/__init__.py
blob: 978370553a1fd23e1687791c3d2b0f6902c03d50 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
from json import JSONDecoder, JSONEncoder

from .helpers import bulk_of_jsons
from ..helpers import nativestr, delist
from .commands import JSONCommands


class JSON(JSONCommands):
    """
    Create a client for talking to json.

    :param decoder:
    :type json.JSONDecoder: An instance of json.JSONDecoder

    :param encoder:
    :type json.JSONEncoder: An instance of json.JSONEncoder
    """

    def __init__(
        self,
        client,
        version=None,
        decoder=JSONDecoder(),
        encoder=JSONEncoder(),
    ):
        """
        Create a client for talking to json.

        :param decoder:
        :type json.JSONDecoder: An instance of json.JSONDecoder

        :param encoder:
        :type json.JSONEncoder: An instance of json.JSONEncoder
        """
        # Set the module commands' callbacks
        self.MODULE_CALLBACKS = {
            "JSON.CLEAR": int,
            "JSON.DEL": int,
            "JSON.FORGET": int,
            "JSON.GET": self._decode,
            "JSON.MGET": bulk_of_jsons(self._decode),
            "JSON.SET": lambda r: r and nativestr(r) == "OK",
            "JSON.NUMINCRBY": self._decode,
            "JSON.NUMMULTBY": self._decode,
            "JSON.TOGGLE": lambda b: b == b"true",
            "JSON.STRAPPEND": int,
            "JSON.STRLEN": int,
            "JSON.ARRAPPEND": int,
            "JSON.ARRINDEX": int,
            "JSON.ARRINSERT": int,
            "JSON.ARRLEN": int,
            "JSON.ARRPOP": self._decode,
            "JSON.ARRTRIM": int,
            "JSON.OBJLEN": int,
            "JSON.OBJKEYS": delist,
            # "JSON.RESP": delist,
            "JSON.DEBUG": int,
        }

        self.client = client
        self.execute_command = client.execute_command
        self.MODULE_VERSION = version

        for key, value in self.MODULE_CALLBACKS.items():
            self.client.set_response_callback(key, value)

        self.__encoder__ = encoder
        self.__decoder__ = decoder

    def _decode(self, obj):
        """Get the decoder."""
        if obj is None:
            return obj

        try:
            return self.__decoder__.decode(obj)
        except TypeError:
            return self.__decoder__.decode(obj.decode())

    def _encode(self, obj):
        """Get the encoder."""
        return self.__encoder__.encode(obj)