from typing import Mapping


def validate_schema(assert_,
                    data,
                    schema,
                    optional=None,
                    on_unexpected_cols='raise'):
    """Check if a data frame complies with a schema

    Parameters
    ----------
    data : pandas.DataFrame
        Data frame to test
    schema : list or dict
        List with column names (will only validate names)
        or dict with column names as keys, dtypes as values (will validate
        names and dtypes)
    optional : list, optional
        List of optional column names, no warns nor errors if they appear
    on_unexpected_cols : str, optional
        One of 'warn', 'raise' or None. If 'warn', it will warn on extra
        columns, if 'raise' it will raise an error, if None it will completely
        ignore extra columns
    """
    if on_unexpected_cols not in {'warn', 'raise', None}:
        raise ValueError("'on_unexpected_cols' must be one of 'warn', 'raise' "
                         "or None")

    optional = optional or {}
    cols = set(data.columns)
    expected = set(schema)

    missing = expected - cols
    unexpected = cols - expected - set(optional)

    msg = '(validate_schema) Missing columns {missing}.'.format(
        missing=missing)
    assert_(not missing, msg)

    if on_unexpected_cols is not None:
        msg = ('(validate_schema) Unexpected columns {unexpected}'.format(
            unexpected=unexpected))
        caller = assert_ if on_unexpected_cols == 'raise' else assert_.warn
        caller(not unexpected, msg)

    # if passing a mapping, schema is validated (even for optional columns)
    for schema_to_validate in [schema, optional]:
        if isinstance(schema_to_validate, Mapping):
            # validate column types (as many as you can)
            dtypes = data.dtypes.astype(str).to_dict()

            for name, dtype in dtypes.items():
                expected = schema_to_validate.get(name)

                if expected is not None:
                    msg = (
                        '(validate_schema) Wrong dtype for column "{name}". '
                        'Expected: "{expected}". Got: "{dtype}"'.format(
                            name=name, expected=expected, dtype=dtype))
                    assert_(dtype == expected, msg)

    return assert_