def test_write_json() -> None:
# tests if it runs if no arg given
df = pl.DataFrame({"a": [1, 2, 3]})
assert (df.write_json() ==
'{"columns":[{"name":"a","datatype":"Int64","values":[1,2,3]}]}')
df = pl.DataFrame({"a": [1, 2, 3], "b": ["a", "b", None]})
expected = df
out = df.write_json(row_oriented=True)
assert out == r"""[{"a":1,"b":"a"},{"a":2,"b":"b"},{"a":3,"b":null}]"""
# test round trip
f = io.BytesIO()
f.write(out.encode()) # type: ignore[attr-defined]
f.seek(0)
df = pl.read_json(f, json_lines=False)
assert df.frame_equal(expected)
out = df.write_json(json_lines=True)
assert (out == r"""{"a":1,"b":"a"}
{"a":2,"b":"b"}
{"a":3,"b":null}
""")
# test round trip
f = io.BytesIO()
f.write(out.encode()) # type: ignore[attr-defined]
f.seek(0)
df = pl.read_json(f, json_lines=True)
expected = pl.DataFrame({"a": [1, 2, 3], "b": ["a", "b", None]})
assert df.frame_equal(expected)
def test_write_json2(df: pl.DataFrame) -> None:
# text-based conversion loses time info
df = df.select(pl.all().exclude(["cat", "time"]))
s = df.write_json(to_string=True)
f = io.BytesIO()
f.write(s.encode())
f.seek(0)
out = pl.read_json(f)
assert df.frame_equal(out, null_equal=True)
file = io.BytesIO()
df.write_json(file)
file.seek(0)
out = pl.read_json(file)
assert df.frame_equal(out, null_equal=True)
def test_to_from_buffer(df: pl.DataFrame) -> None:
for buf in (io.BytesIO(), io.StringIO()):
df.write_json(buf)
buf.seek(0)
read_df = pl.read_json(buf)
read_df = read_df.with_columns(
[pl.col("cat").cast(pl.Categorical),
pl.col("time").cast(pl.Time)])
assert_frame_equal_local_categoricals(df, read_df)
def test_to_json(df):
s = df.to_json(to_string=True)
out = pl.read_json(s)
assert df.frame_equal(out, null_equal=True)
def test_to_from_file(io_test_dir: str, df: pl.DataFrame) -> None:
f = os.path.join(io_test_dir, "small.json")
df.write_json(f)
out = pl.read_json(f)
assert out.frame_equal(df)