Skip to content

Commit

Permalink
Use Arrows for_json for JSON output
Browse files Browse the repository at this point in the history
Arrow provides a `for_json()` function encoding Arrow objects for JSON
output. Using the default-parameter of `json.dumps()`, this function is
invoked on Arrow objects which can then be encoded for JSON output.

Reference PR: #330 

Fix #329
  • Loading branch information
rbialon authored and jmaupetit committed Oct 16, 2019
1 parent 4a40152 commit f46a6ae
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 1 deletion.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
15 changes: 15 additions & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
safe_save,
parse_tags,
PY2,
json_arrow_encoder,
)
from . import mock_datetime

Expand Down Expand Up @@ -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()
4 changes: 3 additions & 1 deletion watson/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
sorted_groupby,
style,
parse_tags,
json_arrow_encoder,
)


Expand Down Expand Up @@ -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)))
Expand Down
18 changes: 18 additions & 0 deletions watson/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import shutil
import sys
import tempfile

try:
from StringIO import StringIO
except ImportError:
Expand All @@ -22,6 +23,8 @@

PY2 = sys.version_info[0] == 2

if not PY2:
from builtins import TypeError, isinstance

try:
text_type = (str, unicode)
Expand Down Expand Up @@ -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))

0 comments on commit f46a6ae

Please sign in to comment.