encode pandas.DataFrame object to jsonable dict object

jsonable_encoder_pandas_dataframe.py

import json
from typing import Optional
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pandas import DataFrame


class JsonableEncoderError(Exception):
    pass


def jsonable_encoder_df(df: 'DataFrame', *,
                        max_row: Optional[int] = None,
                        without_data: bool = False,
                        without_type: bool = False,
                        without_index: bool = True) -> dict:
    """encode DataFrame object to jsonable dict object

    Args:
        df(DataFrame): dataframe object to encode.
        max_row(int): [Optional] max rows to keep, default None.
        without_data(bool): whether drop rows data or not, default False(keep).
        without_type(bool): whether drop types of column or not, default False(keep).
        without_index(bool): whether drop index of dataframe or not, default True(drop).

    Returns:
       dict: jsonable data

    Raises:
        JsonableEncoderError: fail to encode DataFrame object to jsonable dict object.
    """
    if max_row:
       df = df.head(max_row)

    try:
        serialized = df.to_json(orient='split', date_format='iso')
        jsonable_ = json.loads(serialized)
    except Exception as exc:
        raise JsonableEncoderError('jsonable encoder dataframe fail.') from exc

    if without_index:
        del jsonable_['index']
    if without_data:
        del jsonable_['data']

    if not without_type:
        try:
            from pandas.io.json import build_table_schema
            types = [d['type'] for d in build_table_schema(df)['fields']
                     if d['name'] != (df.index.name or 'index')]
        except Exception as exc:
            raise JsonableEncoderError('infer dataframe types failed.') from exc
        jsonable_['column_types'] = types
    return jsonable_