forked from p15670423/monkey
Island: Refactor representations.py to extend JSON encoder
This commit is contained in:
parent
3958ad3e92
commit
4bbcffabd3
|
@ -2,46 +2,31 @@ from datetime import datetime
|
|||
from enum import Enum
|
||||
|
||||
import bson
|
||||
from bson.json_util import dumps
|
||||
from flask import make_response
|
||||
from flask.json import JSONEncoder, dumps
|
||||
|
||||
from common.utils import IJSONSerializable
|
||||
|
||||
|
||||
def _normalize_obj(obj):
|
||||
if ("_id" in obj) and ("id" not in obj):
|
||||
obj["id"] = obj["_id"]
|
||||
del obj["_id"]
|
||||
|
||||
for key, value in list(obj.items()):
|
||||
obj[key] = _normalize_value(value)
|
||||
return obj
|
||||
|
||||
|
||||
def _normalize_value(value):
|
||||
class APIEncoder(JSONEncoder):
|
||||
def default(self, value):
|
||||
# ObjectId is serializible by default, but returns a dict
|
||||
# So serialize it first into a plain string
|
||||
if isinstance(value, bson.objectid.ObjectId):
|
||||
return str(value)
|
||||
|
||||
if isinstance(value, list):
|
||||
return [_normalize_value(_value) for _value in value]
|
||||
if isinstance(value, tuple):
|
||||
return tuple((_normalize_value(_value) for _value in value))
|
||||
if type(value) == dict:
|
||||
return _normalize_obj(value)
|
||||
if isinstance(value, datetime):
|
||||
return str(value)
|
||||
if issubclass(type(value), Enum):
|
||||
return value.name
|
||||
|
||||
if issubclass(type(value), IJSONSerializable):
|
||||
return value.__class__.to_json(value)
|
||||
try:
|
||||
dumps(value)
|
||||
return value
|
||||
return JSONEncoder.default(self, value)
|
||||
except TypeError:
|
||||
return value.__dict__
|
||||
|
||||
|
||||
def output_json(value, code, headers=None):
|
||||
value = _normalize_value(value)
|
||||
resp = make_response(dumps(value), code)
|
||||
resp = make_response(dumps(value, APIEncoder), code)
|
||||
resp.headers.extend(headers or {})
|
||||
return resp
|
||||
|
|
|
@ -1,10 +1,11 @@
|
|||
import json
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
import bson
|
||||
|
||||
from monkey_island.cc.services.representations import _normalize_value
|
||||
from monkey_island.cc.services.representations import APIEncoder
|
||||
|
||||
|
||||
@dataclass
|
||||
|
@ -17,22 +18,24 @@ bogus_object1 = MockClass(1)
|
|||
bogus_object2 = MockClass(2)
|
||||
|
||||
|
||||
def test_normalize_dicts():
|
||||
assert {} == _normalize_value({})
|
||||
def test_api_encoder_dicts():
|
||||
assert json.dumps({}) == json.dumps({}, cls=APIEncoder)
|
||||
|
||||
assert {"a": "a"} == _normalize_value({"a": "a"})
|
||||
assert json.dumps({"a": "a"}) == json.dumps({"a": "a"}, cls=APIEncoder)
|
||||
|
||||
assert {"id": 12345} == _normalize_value({"id": 12345})
|
||||
assert json.dumps({"id": 12345}) == json.dumps({"id": 12345}, cls=APIEncoder)
|
||||
|
||||
assert {"id": obj_id_str} == _normalize_value({"id": bson.objectid.ObjectId(obj_id_str)})
|
||||
assert json.dumps({"id": obj_id_str}) == json.dumps(
|
||||
{"id": bson.objectid.ObjectId(obj_id_str)}, cls=APIEncoder
|
||||
)
|
||||
|
||||
dt = datetime.now()
|
||||
expected = {"a": str(dt)}
|
||||
result = _normalize_value({"a": dt})
|
||||
assert expected == result
|
||||
result = json.dumps({"a": dt}, cls=APIEncoder)
|
||||
assert json.dumps(expected) == result
|
||||
|
||||
|
||||
def test_normalize_complex():
|
||||
def test_api_encoder_complex():
|
||||
bogus_dict = {
|
||||
"a": [
|
||||
{
|
||||
|
@ -44,26 +47,28 @@ def test_normalize_complex():
|
|||
}
|
||||
|
||||
expected_dict = {"a": [{"ba": obj_id_str, "bb": obj_id_str}], "b": {"id": obj_id_str}}
|
||||
assert expected_dict == _normalize_value(bogus_dict)
|
||||
assert json.dumps(expected_dict) == json.dumps(bogus_dict, cls=APIEncoder)
|
||||
|
||||
|
||||
def test_normalize_list():
|
||||
def test_api_encoder_list():
|
||||
bogus_list = [bson.objectid.ObjectId(obj_id_str), {"a": "b"}, {"object": [bogus_object1]}]
|
||||
|
||||
expected_list = [obj_id_str, {"a": "b"}, {"object": [{"a": 1}]}]
|
||||
assert expected_list == _normalize_value(bogus_list)
|
||||
assert json.dumps(expected_list) == json.dumps(bogus_list, cls=APIEncoder)
|
||||
|
||||
|
||||
def test_normalize_enum():
|
||||
def test_api_encoder_enum():
|
||||
class BogusEnum(Enum):
|
||||
bogus_val = "Bogus"
|
||||
|
||||
my_obj = {"something": "something", "my_enum": BogusEnum.bogus_val}
|
||||
|
||||
assert {"something": "something", "my_enum": "bogus_val"} == _normalize_value(my_obj)
|
||||
assert json.dumps({"something": "something", "my_enum": "bogus_val"}) == json.dumps(
|
||||
my_obj, cls=APIEncoder
|
||||
)
|
||||
|
||||
|
||||
def test_normalize_tuple():
|
||||
bogus_tuple = [{"my_tuple": (bogus_object1, bogus_object2, b"one_two")}]
|
||||
expected_tuple = [{"my_tuple": ({"a": 1}, {"a": 2}, b"one_two")}]
|
||||
assert expected_tuple == _normalize_value(bogus_tuple)
|
||||
def test_api_encoder_tuple():
|
||||
bogus_tuple = [{"my_tuple": (bogus_object1, bogus_object2, "string")}]
|
||||
expected_tuple = [{"my_tuple": ({"a": 1}, {"a": 2}, "string")}]
|
||||
assert json.dumps(expected_tuple) == json.dumps(bogus_tuple, cls=APIEncoder)
|
||||
|
|
Loading…
Reference in New Issue