Serialize and deserialize polars datatypes to json

polars_datatypes_json.py

import json
import polars as pl

def dtype_to_json(dtype: pl.DataType) -> str:
    return json.dumps(str(dtype))


def json_to_dtype(json_dtype_str: str) -> pl.DataType:
    from polars.datatypes.classes import (  # noqa F401
        Array,
        Binary,
        Boolean,
        Categorical,
        Date,
        Datetime,
        Decimal,
        Duration,
        Enum,
        Float32,
        Float64,
        Int8,
        Int16,
        Int32,
        Int64,
        List,
        Null,
        Object,
        Struct,
        Time,
        UInt8,
        UInt16,
        UInt32,
        UInt64,
        Unknown,
        Utf8,
    )

    return eval(json.loads(json_dtype_str))

# tests

dtypes = [
    pl.Datetime(time_zone="UTC"),
    pl.List(pl.Array(pl.Float32(), 3)),
    pl.Float32(),
    pl.Struct({"a": pl.Int8(), "b": pl.List(pl.Utf8())}),
]

for dtype in dtypes:
    assert dtype == json_to_dtype(dtype_to_json(dtype))