diff --git a/src/fsql/api.py b/src/fsql/api.py index ece31c4..528fcbd 100644 --- a/src/fsql/api.py +++ b/src/fsql/api.py @@ -147,6 +147,9 @@ def write_object( elif format == "csv": with fs.open(url_suff, "wb") as fd: data.to_csv(fd) + elif format == "json": + with fs.open(url_suff, "wb") as fd: + data.to_json(fd) else: raise ValueError(f"unsupported format for dataframe writing: {format}") elif isinstance(data, io.StringIO) or isinstance(data, io.BytesIO): diff --git a/tests/conftest.py b/tests/conftest.py index 6af441e..cf75ead 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ """Sets a mock/fake of the S3 filesystem for any `fsql`-based usage.""" +import json import os import fsspec @@ -14,6 +15,10 @@ def put_s3_file(self, data, url): with self.s3fs.open(url, "wb") as fd: fd.write(data) + def read_json_file(self, url): + with self.s3fs.open(url, "r") as fd: + return json.load(fd) + @pytest.fixture def helper(): diff --git a/tests/test_write_object.py b/tests/test_write_object.py index b78b702..77057b6 100644 --- a/tests/test_write_object.py +++ b/tests/test_write_object.py @@ -1,4 +1,5 @@ import io +import json import pandas as pd import pytest @@ -87,3 +88,28 @@ def test_write_vanilla_bytes(tmpdir): with open(path, "rb") as f: extracted = f.read() assert extracted == data + + +def test_write_json_s3(helper): + """Writes a dataframe as json, tests that read works.""" + bucket = "test-bouquet" + fs = helper.s3fs + fs.mkdir(bucket) + + input = pd.DataFrame({"k1": [1, 2], "k2": [3, 4]}, index=["one", "two"]) + write_object("s3://test-bouquet/my_df.json", input, format="json") + output = helper.read_json_file("s3://test-bouquet/my_df.json") + expected_output = {"k1": {"one": 1, "two": 2}, "k2": {"one": 3, "two": 4}} + assert output == expected_output + + +def test_write_json(tmpdir): + """Writes a json file, tests that read works.""" + input = pd.DataFrame({"k1": [1, 2], "k2": [3, 4]}, index=["one", "two"]) + path_base = tmpdir.join("my_file.json") + url = f"file://{path_base}" + write_object(url, input, format="json") + with open(path_base, "r") as f: + output = json.load(f) + expected_output = {"k1": {"one": 1, "two": 2}, "k2": {"one": 3, "two": 4}} + assert output == expected_output