Skip to content

Commit

Permalink
Add ContextBoolEncoder to facilitate json serialization of ContextBools.
Browse files Browse the repository at this point in the history
  • Loading branch information
DanSava committed Jan 25, 2024
1 parent 797b49a commit b5bf251
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 2 deletions.
8 changes: 7 additions & 1 deletion src/ert/config/parsing/context_values.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,16 @@
from typing import List, TypeVar, Union, no_type_check
from json import JSONEncoder
from typing import Any, List, TypeVar, Union, no_type_check

from .file_context_token import FileContextToken

# mypy: disable-error-code="attr-defined"


class ContextBoolEncoder(JSONEncoder):
def default(self, o: Any) -> Any:
return o.val if isinstance(o, ContextBool) else JSONEncoder.default(self, o)


class ContextBool:
def __init__(
self, val: bool, token: FileContextToken, keyword_token: FileContextToken
Expand Down
3 changes: 2 additions & 1 deletion src/ert/storage/local_experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
SummaryConfig,
SurfaceConfig,
)
from ert.config.parsing.context_values import ContextBoolEncoder
from ert.config.response_config import ResponseConfig

if TYPE_CHECKING:
Expand Down Expand Up @@ -230,4 +231,4 @@ def write_simulation_arguments(
with open(
self.mount_point / self._simulation_arguments_file, "w", encoding="utf-8"
) as f:
json.dump(dataclasses.asdict(info), f)
json.dump(dataclasses.asdict(info), f, cls=ContextBoolEncoder)
43 changes: 43 additions & 0 deletions tests/unit_tests/config/test_ert_config.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import json
import logging
import os
import os.path
Expand All @@ -19,6 +20,14 @@
)
from ert.config.ert_config import site_config_location
from ert.config.parsing import ConfigKeys, ConfigWarning
from ert.config.parsing.context_values import (
ContextBool,
ContextBoolEncoder,
ContextFloat,
ContextInt,
ContextList,
ContextString,
)
from ert.job_queue import Driver

from .config_dict_generator import config_generators
Expand Down Expand Up @@ -1545,3 +1554,37 @@ def test_that_removed_analysis_module_keywords_raises_error(
match=error_msg,
):
_ = ErtConfig.from_file(test_config_file_name)


@pytest.mark.usefixtures("use_tmpdir")
def test_that_context_types_are_json_serializable():
bf = ContextBool(val=False, token=None, keyword_token=None)
bt = ContextBool(val=True, token=None, keyword_token=None)
i = ContextInt(val=23, token=None, keyword_token=None)
s = ContextString(val="forty_two", token=None, keyword_token=None)
fl = ContextFloat(val=4.2, token=None, keyword_token=None)
cl = ContextList.with_values(None, values=[bf, bt, i, s, fl])

payload = {
"context_bool_false": bf,
"context_bool_true": bt,
"context_int": i,
"context_str": s,
"context_float": fl,
"context_list": cl,
}

with open("test.json", "w", encoding="utf-8") as f:
json.dump(payload, f, cls=ContextBoolEncoder)

with open("test.json", "r", encoding="utf-8") as f:
r = json.load(f)

assert isinstance(r["context_bool_false"], bool)
assert isinstance(r["context_bool_true"], bool)
assert r["context_bool_false"] is False
assert r["context_bool_true"] is True
assert isinstance(r["context_int"], int)
assert isinstance(r["context_str"], str)
assert isinstance(r["context_float"], float)
assert isinstance(r["context_list"], list)

0 comments on commit b5bf251

Please sign in to comment.