From 74328caecaacfddacfee5eaa9888e30a2e6f4e8e Mon Sep 17 00:00:00 2001 From: Colton Hicks Date: Fri, 24 Mar 2023 11:38:46 -0700 Subject: [PATCH] Added ability for .json() to accept kwargs. model.json() can now create a pretty string with .json(indent=4) --- qcelemental/models/basemodels.py | 19 +++++++++++-------- qcelemental/tests/test_model_results.py | 8 ++++++++ qcelemental/util/serialization.py | 12 ++++++++---- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/qcelemental/models/basemodels.py b/qcelemental/models/basemodels.py index d0edfe61..d7da1727 100644 --- a/qcelemental/models/basemodels.py +++ b/qcelemental/models/basemodels.py @@ -127,6 +127,7 @@ def serialize( exclude_unset: Optional[bool] = None, exclude_defaults: Optional[bool] = None, exclude_none: Optional[bool] = None, + **kwargs, ) -> Union[bytes, str]: r"""Generates a serialized representation of the model @@ -144,6 +145,8 @@ def serialize( If True, skips fields that have set or defaulted values equal to the default. exclude_none If True, skips fields that have value ``None``. + kwargs + For passing additional values like indent to serialization functions Returns ------- @@ -151,21 +154,21 @@ def serialize( The serialized model. """ - kwargs = {} + dict_kwargs = {} if include: - kwargs["include"] = include + dict_kwargs["include"] = include if exclude: - kwargs["exclude"] = exclude + dict_kwargs["exclude"] = exclude if exclude_unset: - kwargs["exclude_unset"] = exclude_unset + dict_kwargs["exclude_unset"] = exclude_unset if exclude_defaults: - kwargs["exclude_defaults"] = exclude_defaults + dict_kwargs["exclude_defaults"] = exclude_defaults if exclude_none: - kwargs["exclude_none"] = exclude_none + dict_kwargs["exclude_none"] = exclude_none - data = self.dict(**kwargs) + data = self.dict(**dict_kwargs) - return serialize(data, encoding=encoding) + return serialize(data, encoding=encoding, **kwargs) def json(self, **kwargs): # Alias JSON here from BaseModel to reflect dict changes diff --git a/qcelemental/tests/test_model_results.py b/qcelemental/tests/test_model_results.py index b77dfdd2..a9725d25 100644 --- a/qcelemental/tests/test_model_results.py +++ b/qcelemental/tests/test_model_results.py @@ -464,6 +464,14 @@ def test_failed_operation(result_data_fixture, request): assert "its all good" in failed_json +def test_model_json_allows_kwargs(result_data_fixture): + result = qcel.models.AtomicResult(**result_data_fixture) + no_indent = result.json() + assert "\n" not in no_indent + indent = result.json(indent=2) + assert "\n" in indent + + def test_result_properties_array(request): lquad = [1, 2, 3, 2, 4, 5, 3, 5, 6] diff --git a/qcelemental/util/serialization.py b/qcelemental/util/serialization.py index a8f719e1..8fc97d27 100644 --- a/qcelemental/util/serialization.py +++ b/qcelemental/util/serialization.py @@ -203,13 +203,15 @@ def default(self, obj: Any) -> Any: return json.JSONEncoder.default(self, obj) -def json_dumps(data: Any) -> str: +def json_dumps(data: Any, **kwargs) -> str: r"""Safe serialization of a Python dictionary to JSON string representation using all known encoders. Parameters ---------- data : Any A encodable python object. + kwargs : Any + Keyword arguments for json.dumps() Returns ------- @@ -217,7 +219,7 @@ def json_dumps(data: Any) -> str: A JSON representation of the data. """ - return json.dumps(data, cls=JSONArrayEncoder) + return json.dumps(data, cls=JSONArrayEncoder, **kwargs) def json_loads(data: str) -> Any: @@ -313,7 +315,7 @@ def msgpack_loads(data: str) -> Any: ## Helper functions -def serialize(data: Any, encoding: str) -> Union[str, bytes]: +def serialize(data: Any, encoding: str, **kwargs) -> Union[str, bytes]: r"""Encoding Python objects using the provided encoder. Parameters @@ -322,6 +324,8 @@ def serialize(data: Any, encoding: str) -> Union[str, bytes]: A encodable python object. encoding : str The type of encoding to perform: {'json', 'json-ext', 'msgpack-ext'} + kwargs: Any + For passing additional kwargs to serialization functions, such as indent Returns ------- @@ -330,7 +334,7 @@ def serialize(data: Any, encoding: str) -> Union[str, bytes]: """ if encoding.lower() == "json": - return json_dumps(data) + return json_dumps(data, **kwargs) elif encoding.lower() == "json-ext": return jsonext_dumps(data) elif encoding.lower() == "msgpack":