diff --git a/CHANGELOG.md b/CHANGELOG.md index 1068d338..4f917810 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Stylize prompt to create new project or tag (#310). - Aggregate calculates wrong time if used with `--current` (#293) - The `start` command now correctly checks if project is empty (#322) +- The `report` and `aggregate` commands with `--json` option now correctly + encode Arrow objects (#329) ## [1.8.0] - 2019-08-26 diff --git a/tests/test_utils.py b/tests/test_utils.py index f1fcdb88..4b84b974 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -34,6 +34,7 @@ safe_save, parse_tags, PY2, + json_arrow_encoder, ) from . import mock_datetime @@ -339,3 +340,17 @@ def test_flatten_report_for_csv(watson): assert result[2]['project'] == 'foo' assert result[2]['tag'] == 'B' assert result[2]['time'] == (4 + 2) * 3600 + + +def test_json_arrow_encoder(): + with pytest.raises(TypeError): + json_arrow_encoder(0) + + with pytest.raises(TypeError): + json_arrow_encoder('foo') + + with pytest.raises(TypeError): + json_arrow_encoder(None) + + now = arrow.utcnow() + assert json_arrow_encoder(now) == now.for_json() diff --git a/watson/cli.py b/watson/cli.py index fee1e907..33b0f16d 100644 --- a/watson/cli.py +++ b/watson/cli.py @@ -31,6 +31,7 @@ sorted_groupby, style, parse_tags, + json_arrow_encoder, ) @@ -605,7 +606,8 @@ def report(watson, current, from_, to, projects, tags, ignore_projects, luna=luna, all=all) if 'json' in output_format and not aggregated: - click.echo(json.dumps(report, indent=4, sort_keys=True)) + click.echo(json.dumps(report, indent=4, sort_keys=True, + default=json_arrow_encoder)) return elif 'csv' in output_format and not aggregated: click.echo(build_csv(flatten_report_for_csv(report))) diff --git a/watson/utils.py b/watson/utils.py index 82a50d66..f90574cf 100644 --- a/watson/utils.py +++ b/watson/utils.py @@ -8,6 +8,7 @@ import shutil import sys import tempfile + try: from StringIO import StringIO except ImportError: @@ -22,6 +23,8 @@ PY2 = sys.version_info[0] == 2 +if not PY2: + from builtins import TypeError, isinstance try: text_type = (str, unicode) @@ -397,3 +400,18 @@ def flatten_report_for_csv(report): 'time': tag['time'] }) return result + + +def json_arrow_encoder(obj): + """ + Encodes Arrow objects for JSON output. + This function can be used with + `json.dumps(..., default=json_arrow_encoder)`, for example. + If the object is not an Arrow type, a TypeError is raised + :param obj: Object to encode + :return: JSON representation of Arrow object as defined by Arrow + """ + if isinstance(obj, arrow.Arrow): + return obj.for_json() + + raise TypeError("Object {} is not JSON serializable".format(obj))