diff --git a/smartnoise_synth_logger/deserialise.py b/smartnoise_synth_logger/deserialise.py index 4dec0f6..10f81ae 100644 --- a/smartnoise_synth_logger/deserialise.py +++ b/smartnoise_synth_logger/deserialise.py @@ -82,7 +82,9 @@ def deserialise_constraints(constraints_json: str) -> dict: deserialised = {} for key, val in json_body[JsonBodyKey.CONSTRAINTS].items(): - if isinstance(val[JsonBodyKey.PARAM], list): + if isinstance(val, str): + deserialised[key] = val + elif isinstance(val[JsonBodyKey.PARAM], list): tranformer_list = [] for t in val[JsonBodyKey.PARAM]: tranformer_list.append( diff --git a/smartnoise_synth_logger/serialise.py b/smartnoise_synth_logger/serialise.py index 4c687af..0b0048a 100644 --- a/smartnoise_synth_logger/serialise.py +++ b/smartnoise_synth_logger/serialise.py @@ -54,23 +54,28 @@ def serialise_constraints(constraints: dict) -> str: } for col_name, col_constraints in constraints.items(): - operator_name = col_constraints.__class__.__name__ + if isinstance(col_constraints, str): + json_body[JsonBodyKey.CONSTRAINTS][col_name] = col_constraints + else: + operator_name = col_constraints.__class__.__name__ - if operator_name == Transformers.CHAIN: - transformer_dict = handle_chain_transformer(col_constraints) - elif operator_name == Transformers.ANONIMIZATION: - transformer_dict = { - JsonBodyKey.TYPE: SSYNTH_TRANSFORMER - + Transformers.ANONIMIZATION, - JsonBodyKey.PARAM: {ANON_PARAM: col_constraints.fake.__name__}, - } - else: # default - transformer_dict = { - JsonBodyKey.TYPE: SSYNTH_TRANSFORMER - + col_constraints.__class__.__name__, - JsonBodyKey.PARAM: get_filtered_params(col_constraints), - } + if operator_name == Transformers.CHAIN: + transformer_dict = handle_chain_transformer(col_constraints) + elif operator_name == Transformers.ANONIMIZATION: + transformer_dict = { + JsonBodyKey.TYPE: SSYNTH_TRANSFORMER + + Transformers.ANONIMIZATION, + JsonBodyKey.PARAM: { + ANON_PARAM: col_constraints.fake.__name__ + }, + } + else: # default + transformer_dict = { + JsonBodyKey.TYPE: SSYNTH_TRANSFORMER + + col_constraints.__class__.__name__, + JsonBodyKey.PARAM: get_filtered_params(col_constraints), + } - json_body[JsonBodyKey.CONSTRAINTS][col_name] = transformer_dict + json_body[JsonBodyKey.CONSTRAINTS][col_name] = transformer_dict return json.dumps(json_body) diff --git a/tests/test_de_serialisation.py b/tests/test_de_serialisation.py index eee5cc8..d2d2b29 100644 --- a/tests/test_de_serialisation.py +++ b/tests/test_de_serialisation.py @@ -29,6 +29,17 @@ def test_anon_serialize(): assert result_json == expected_json_updated +def test_anon_str_serialize(): + example_constraints = {"id": "email"} + result_json = serialise_constraints(example_constraints) + + expected_json = """{"module": "smartnoise-synth", "version": "1.0.4", "constraints": {"id": "email"}}""" # noqa + expected_json_updated = expected_json.replace( + "1.0.4", pkg_resources.get_distribution(SSYNTH).version + ) + assert result_json == expected_json_updated + + def test_chain_serialize(): example_constraints = { "income": ChainTransformer(