pandas testing

pandas_testing.py

from StringIO import StringIO
from itertools import islice, izip_longest, tee

import pandas as pd
from pandas.util.testing import assert_frame_equal

def read_data_frame_str(s, first_col_index=False):
    """Read a markdown-style description of a table into a DataFrame.
    :param str s:
    :param bool first_col_index: if `True`, use the first column as an index.
    :rtype: pandas.DataFrame
    """
    df = pd.read_csv(
        StringIO(s),
        index_col=None,
        sep=r' +\| +',
        skiprows=[0, 2],
        comment='  #',
        engine='python',
    )
    if first_col_index:
        df = df.set_index(df.columns[0])
    return df


def read_table_str(session, table, s):
    """Read the string description of a table and INSERT it into a table."""
    df = read_data_frame_str(s, first_col_index=True)
    df.to_sql(table.__tablename__, session.connection(), if_exists='append')


def assert_equal(actual, expected):
    """Assert 2 DataFrames are equal enough for our purposes."""
    assert_frame_equal(
        actual.reset_index().drop('index', axis=1, errors='ignore'),
        expected.reset_index().drop('index', axis=1, errors='ignore'),
        check_dtype=False,
        check_column_type=False,
    )