opdb/user_code/precompiled/__serializer__.py

import array
from collections.abc import Iterable

import orjson
import ujson as json

from .listnode import ListNode
from .nestedinteger import NestedInteger
from .treenode import TreeNode


class __Serializer__:
    def _serialize_int(self, x):
        return str(x)

    # TODO: precision
    # if x = 3.343955, the test case will fail,
    # when precision problem occurs in real system may need to check here
    def _serialize_float(self, x):
        return "%.5f" % x

    def _serialize_str(self, x):
        return json.dumps(x, escape_forward_slashes=False)

    def _serialize_bool(self, x):
        return "true" if x else "false"

    # TODO: depreciated. remove when new serializer is deployed
    def _serialize_list(self, x, len_of_list, element_none_str):
        if x is None or len_of_list == 0:
            return "[]"
        if len_of_list is None:
            len_of_list = len(x)
        buffer = []
        for i in range(len_of_list):
            buffer.append(
                "".join(self.serialize(x[i], none_str=element_none_str))
            )
        return "[%s]" % ",".join(buffer)

    def serialize_list(self, x, t):
        if x is None:
            return "[]"
        return "[" + ",".join([self._serialize(e, t) for e in x]) + "]"

    # TODO: depreciated. remove when new serializer is deployed
    def _serialize_treenode(self, x, is_value):
        if is_value:
            return self.serialize(x.val) if x else "null"
        else:
            return TreeNode.serialize(x)

    def _serialize(self, x, t):
        if t[-2:] == "[]":
            return self.serialize_list(x, t[:-2])
        elif t[-1:] == ">":
            return self.serialize_list(x, t[5:-1])
        elif t == "integer":
            return self._serialize_int(x)
        elif t == "long":
            return self._serialize_int(x)
        elif t == "double":
            return self._serialize_float(x)
        elif t == "character":
            return self._serialize_str(x)
        elif t == "boolean":
            return self._serialize_bool(x)
        elif t == "string":
            return self._serialize_str(x)
        elif t == "ListNode":
            return ListNode.serialize(x)
        elif t == "TreeNode":
            return TreeNode.serialize(x)
        elif t == "NestedInteger":
            return NestedInteger._serialize(x, self)
        else:
            raise Exception("Type %s: Not implemented" % t)

    # TODO: depreciated. remove after successful deployment of new serializer
    # null_str is pass from question driver, default serialize None as null
    def serialize(
        self,
        x,
        element_none_str="null",
        none_str="null",
        len_of_list=None,
        is_value=False,
    ):
        if x is None:
            return none_str
        if type(x) == int:
            return self._serialize_int(x)
        elif type(x) == float:
            return self._serialize_float(x)
        elif type(x) == str:
            return self._serialize_str(x)
        elif type(x) == bool:
            return self._serialize_bool(x)
        elif isinstance(x, array.array):
            return self._serialize_list(
                x.tolist(), len_of_list, element_none_str
            )
        elif isinstance(x, list):
            return self._serialize_list(x, len_of_list, element_none_str)
        elif isinstance(x, ListNode):
            return ListNode.serialize(x)
        elif isinstance(x, TreeNode):
            return self._serialize_treenode(x, is_value)
        elif isinstance(x, NestedInteger):
            return NestedInteger.serialize(x, self)
        else:
            raise Exception("Type %s: Not implemented" % str(type(x)))


def serializer_node(obj):
    """
    如果node的元素为None，返回的是空列表
    """
    if isinstance(obj, ListNode):
        return obj._list_node_to_array()
    elif isinstance(obj, TreeNode):
        return obj._tree_node_to_array()
    elif isinstance(obj, NestedInteger):
        return obj._nested_integer_to_token()
    elif obj is None:
        return []
    else:
        raise Exception("Type %s cannot be serialized" % str(type(obj)))


def check_type(type_str):
    while len(type_str):
        if type_str[-2:] == "[]":
            type_str = type_str[:-2]
        elif type_str[-1:] == ">":
            type_str = type_str[5:-1]
        elif type_str in (
            "integer",
            "long",
            "double",
            "character",
            "boolean",
            "string",
            "ListNode",
            "TreeNode",
            "NestedInteger",
        ):
            type_str = ""
        else:
            return False
    return True


class __SerializerRapid__:
    def _serialize_float_or_float_list(self, obj):
        """
        double 类型特殊处理:
            保留5位
        """
        if not isinstance(obj, Iterable):
            return "%.5f" % obj
        float_list = [self._serialize_float_or_float_list(i) for i in obj]
        return "[%s]" % ",".join(float_list)

    def _serialize_default(self, obj, type_str):
        if type_str[-2:] == "[]":
            return [self._serialize_default(x, type_str[:-2]) for x in obj]
        elif type_str[-1:] == ">":
            return [self._serialize_default(x, type_str[5:-1]) for x in obj]
        else:
            return serializer_node(obj)

    def _serialize(self, obj, type_str):
        """
        注意： 这里ListNode, TreeNode不能直接调用orjson的default功能
            当ListNode元素值为None, 需要返回空列表[], 而不是null
        """
        if not check_type(type_str):
            raise Exception("Type %s: Not implemented" % type_str)

        if "double" in type_str:
            return self._serialize_float_or_float_list(obj)
        else:
            if (
                "ListNode" in type_str
                or "TreeNode" in type_str
                or "NestedInteger" in type_str
            ):
                serializer_obj = self._serialize_default(obj, type_str)
            else:
                serializer_obj = obj
            return bytes.decode(orjson.dumps(serializer_obj))