diff --git a/flow/record/jsonpacker.py b/flow/record/jsonpacker.py index 344b5538..d971b818 100644 --- a/flow/record/jsonpacker.py +++ b/flow/record/jsonpacker.py @@ -49,12 +49,7 @@ def pack_obj(self, obj: Any) -> dict | str: serial["_type"] = "record" serial["_recorddescriptor"] = obj._desc.identifier - for field_type, field_name in obj._desc.get_field_tuples(): - # Boolean field types should be cast to a bool instead of staying ints - if field_type == "boolean" and isinstance(serial[field_name], int): - serial[field_name] = bool(serial[field_name]) - - return serial + return self.convert_basic_types(serial) if isinstance(obj, RecordDescriptor): return { "_type": "recorddescriptor", @@ -102,7 +97,19 @@ def unpack_obj(self, obj: Any) -> RecordDescriptor | Record | Any: return RecordDescriptor._unpack(*data) return obj - def pack(self, obj: Record | RecordDescriptor) -> str: + def convert_basic_types(self, obj: Any) -> Any: + """Explicitly convert some basic types when packing to JSON.""" + if isinstance(obj, fieldtypes.boolean): + return bool(obj) + if isinstance(obj, dict): + return {k: self.convert_basic_types(v) for k, v in obj.items()} + if isinstance(obj, list): + return [self.convert_basic_types(item) for item in obj] + return obj + + def pack(self, obj: Record | RecordDescriptor | dict) -> str: + if isinstance(obj, dict): + obj = self.convert_basic_types(obj) return json.dumps(obj, default=self.pack_obj, indent=self.indent) def unpack(self, d: str) -> RecordDescriptor | Record: diff --git a/tests/fieldtypes/test_boolean.py b/tests/fieldtypes/test_boolean.py new file mode 100644 index 00000000..e3ee848f --- /dev/null +++ b/tests/fieldtypes/test_boolean.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import pytest + +from flow.record.base import RecordDescriptor + + +def test_boolean() -> None: + TestRecord = RecordDescriptor( + "test/boolean", + [ + ("boolean", "booltrue"), + ("boolean", "boolfalse"), + ], + ) + + r = TestRecord(True, False) + assert bool(r.booltrue) is True + assert bool(r.boolfalse) is False + + r = TestRecord(1, 0) + assert bool(r.booltrue) is True + assert bool(r.boolfalse) is False + + assert str(r.booltrue) == "True" + assert str(r.boolfalse) == "False" + + assert repr(r.booltrue) == "True" + assert repr(r.boolfalse) == "False" + + with pytest.raises(ValueError, match="Value not a valid boolean value"): + TestRecord(2, -1) + + with pytest.raises(ValueError, match="invalid literal for int"): + TestRecord("True", "False") diff --git a/tests/fieldtypes/test_fieldtypes.py b/tests/fieldtypes/test_fieldtypes.py index 46e9a392..6c6eb099 100644 --- a/tests/fieldtypes/test_fieldtypes.py +++ b/tests/fieldtypes/test_fieldtypes.py @@ -293,36 +293,6 @@ def test_dictlist() -> None: assert r.hits[1]["b"] == 4 -def test_boolean() -> None: - TestRecord = RecordDescriptor( - "test/boolean", - [ - ("boolean", "booltrue"), - ("boolean", "boolfalse"), - ], - ) - - r = TestRecord(True, False) - assert bool(r.booltrue) is True - assert bool(r.boolfalse) is False - - r = TestRecord(1, 0) - assert bool(r.booltrue) is True - assert bool(r.boolfalse) is False - - assert str(r.booltrue) == "True" - assert str(r.boolfalse) == "False" - - assert repr(r.booltrue) == "True" - assert repr(r.boolfalse) == "False" - - with pytest.raises(ValueError, match="Value not a valid boolean value"): - r = TestRecord(2, -1) - - with pytest.raises(ValueError, match="invalid literal for int"): - r = TestRecord("True", "False") - - def test_float() -> None: TestRecord = RecordDescriptor( "test/float", diff --git a/tests/packer/test_json_packer.py b/tests/packer/test_json_packer.py index 025b30d8..a6d94db1 100644 --- a/tests/packer/test_json_packer.py +++ b/tests/packer/test_json_packer.py @@ -93,6 +93,16 @@ def test_record_pack_bool_regression() -> None: # pack the json string back to a record and make sure it is the same as before assert packer.unpack(data) == record + # Make sure the same applies to an OrderedDict, which is how JsonRecordPacker is invoked for + # the Elastic adapter. + rdict = record._asdict() + data = packer.pack(rdict) + assert data.startswith('{"some_varint": 1, "some_uint": 0, "some_boolean": false, ') + + # test that packer.pack has no side effects on rdict + assert rdict == record._asdict() + assert isinstance(rdict["some_boolean"], fieldtypes.boolean) + def test_record_pack_surrogateescape() -> None: TestRecord = RecordDescriptor(