class GroupedDataFrame(pd.DataFrame):
    """
    Grouped DataFrame

    This is just a :class:`pandas.DataFrame` with information
    on how to do the grouping.
    """
    # See: subclassing-pandas-data-structures at
    # http://pandas.pydata.org/pandas-docs/stable/internals.html
    _metadata = ['plydata_groups']

    plydata_groups = None

    def __init__(self, data=None, groups=None, **kwargs):
        super().__init__(data=data, **kwargs)
        if groups is not None:
            self.plydata_groups = list(pd.unique(groups))

    @property
    def _constructor(self):
        return GroupedDataFrame

    def __repr__(self):
        s = 'groups: {}\n{}'.format(
            self.plydata_groups,
            super().__repr__())
        return s

    def groupby(self, by=None, **kwargs):
        """
        Group by and do not sort (unless specified)

        For plydata use cases, there is no need to specify
        group columns.
        """
        if by is None:
            by = self.plydata_groups

        # Turn off sorting by groups messes with some verbs
        if 'sort' not in kwargs:
            kwargs['sort'] = False

        return super().groupby(by, **kwargs)

    def group_indices(self):
        """
        Return group indices
        """
        # No groups
        if not self.plydata_groups:
            return np.ones(len(self), dtype=int)

        grouper = self.groupby()
        indices = np.empty(len(self), dtype=int)
        for i, (_, idx) in enumerate(sorted(grouper.indices.items())):
            indices[idx] = i
        return indices

    def _repr_html_(self, *args, **kwargs):
        cell = '<td>{}</td>'
        td_gen = (cell.format(x) for x in self.plydata_groups)
        group_html = """
        <table border="1" class="dataframe">
            <tbody>
                <tr>
                    <th>groups</th>
                    {}
                </tr>
            </tbody>
        </table>
        """.format(''.join(td_gen))
        return group_html + super()._repr_html_(*args, **kwargs)

    def to_html(self, *args, **kwargs):
        return self._repr_html_(*args, **kwargs)

    def equals(self, other):
        try:
            same_groups = self.plydata_groups == other.plydata_groups
        except AttributeError:
            same_groups = False

        return same_groups and super().equals(other) def get_option(name):
    """
    Get plydata option

    Parameters
    ----------
    name : str
        Name of the option
    """
    if name not in OPTIONS:
        raise ValueError("Unknown option {!r}".format(name))

    return globals()[name] def register_implementations(module, verb_names, datatype):
    """
    Register verb implementations in the module

    Parameters
    ----------
    module : module
        Module with the implementations.
    verb_names : list
        Names of verbs implemented in the module.
    datatype : str
        A name of the datatype implemented. e.g 'dataframe',
        'dict'
    """
    for name in verb_names:
        REGISTRY[datatype][name] = module[name] def Q(name):
    """
    Quote a variable name

    A way to 'quote' variable names, especially ones that do not otherwise
    meet Python's variable name rules.

    Parameters
    ----------
    name : str
        Name of variable

    Returns
    -------
    value : object
        Value of variable

    Examples
    --------
    >>> import pandas as pd
    >>> from plydata import define
    >>> df = pd.DataFrame({'class': [10, 20, 30]})

    Since ``class`` is a reserved python keyword it cannot be a variable
    name, and therefore cannot be used in an expression without quoting it.

    >>> df >> define(y='class+1')
    Traceback (most recent call last):
      File "<string>", line 1
        class+1
            ^
    SyntaxError: invalid syntax

    >>> df >> define(y='Q("class")+1')
       class   y
    0     10  11
    1     20  21
    2     30  31

    Note that it is ``'Q("some name")'`` and not ``'Q(some name)'``.
    As in the above example, you do not need to ``import`` ``Q`` before
    you can use it.
    """
    env = EvalEnvironment.capture(1)
    try:
        return env.namespace[name]
    except KeyError:
        raise NameError("No data named {!r} found".format(name)) def get_empty_env():
    """
    Return an empty environment

    This is for testing or documentation purposes
    """
    return EvalEnvironment(namespaces={}) @contextmanager
def regular_index(*dfs):
    """
    Change & restore the indices of dataframes

    Dataframe with duplicate values can be hard to work with.
    When split and recombined, you cannot restore the row order.
    This can be the case even if the index has unique but
    irregular/unordered. This contextmanager resets the unordered
    indices of any dataframe passed to it, on exit it restores
    the original index.

    A regular index is of the form::

        RangeIndex(start=0, stop=n, step=1)

    Parameters
    ----------
    dfs : tuple
        Dataframes

    Yields
    ------
    dfs : tuple
        Dataframe

    Examples
    --------
    Create dataframes with different indices

    >>> df1 = pd.DataFrame([4, 3, 2, 1])
    >>> df2 = pd.DataFrame([3, 2, 1], index=[3, 0, 0])
    >>> df3 = pd.DataFrame([11, 12, 13], index=[11, 12, 13])

    Within the contexmanager all frames have nice range indices

    >>> with regular_index(df1, df2, df3):
    ...     print(df1.index)
    ...     print(df2.index)
    ...     print(df3.index)
    RangeIndex(start=0, stop=4, step=1)
    RangeIndex(start=0, stop=3, step=1)
    RangeIndex(start=0, stop=3, step=1)

    Indices restored

    >>> df1.index
    RangeIndex(start=0, stop=4, step=1)
    >>> df2.index
    Int64Index([3, 0, 0], dtype='int64')
    >>> df3.index
    Int64Index([11, 12, 13], dtype='int64')
    """
    original_index = [df.index for df in dfs]
    have_bad_index = [not isinstance(df.index, pd.RangeIndex)
                      for df in dfs]

    for df, bad in zip(dfs, have_bad_index):
        if bad:
            df.reset_index(drop=True, inplace=True)

    try:
        yield dfs
    finally:
        for df, bad, idx in zip(dfs, have_bad_index, original_index):
            if bad and len(df.index) == len(idx):
                df.index = idx def unique(lst):
    """
    Return unique elements

    :class:`pandas.unique` and :class:`numpy.unique` cast
    mixed type lists to the same type. They are faster, but
    some times we want to maintain the type.

    Parameters
    ----------
    lst : list-like
        List of items

    Returns
    -------
    out : list
        Unique items in the order that they appear in the
        input.

    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> lst = ['one', 'two', 123, 'three']
    >>> pd.unique(lst)
    array(['one', 'two', '123', 'three'], dtype='<U5')
    >>> np.unique(lst)
    array(['123', 'one', 'three', 'two'],
          dtype='<U5')
    >>> unique(lst)
    ['one', 'two', 123, 'three']

    pandas and numpy cast 123 to a string!, and numpy does not
    even maintain the order.
    """
    seen = set()

    def make_seen(x):
        seen.add(x)
        return x

    return [make_seen(x) for x in lst if x not in seen] class Evaluator:
    """
    Evaluator for expressions of a verb

    The Evaluator is responsible for breaking up the data,
    doing evaluations and putting the pieces together

    Parameters
    ----------
    verb : verb
        Verb with expressions to evaluate
    keep_index : bool
        If True, the evaluation method *tries* to create a
        dataframe with the same index as the original.
        Ultimately whether this succeeds depends on the
        expressions to be evaluated.
    keep_groups : bool
        If True, the resulting data will be grouped in the input
        data was grouped.
    drop : bool
        If True drop unselected columns. (Used by create)
    """
    def __init__(self, verb, keep_index=True, keep_groups=True, drop=False):
        self.data = verb.data
        self.expressions = verb.expressions
        self.env = get_empty_env() if verb.env is None else verb.env
        self.keep_index = keep_index
        self.keep_groups = keep_groups
        self.drop = drop

        try:
            self.groups = verb.data.plydata_groups
        except AttributeError:
            self.groups = None

    def process(self):
        """
        Run the expressions

        Returns
        -------
        out : pandas.DataFrame
            Resulting data
        """
        # Short cut
        if self._all_expressions_evaluated():
            if self.drop:
                # Drop extra columns. They do not correspond to
                # any expressions.
                columns = [expr.column for expr in self.expressions]
                self.data = self.data.loc[:, columns]
            return self.data

        # group_by
        # evaluate expressions
        # combine columns
        # concat evalutated group data and clean up index and group
        gdfs = self._get_group_dataframes()
        egdfs = self._evaluate_expressions(gdfs)
        edata = self._concat(egdfs)
        return edata

    def _all_expressions_evaluated(self):
        """
        Return True if all expressions match with the columns

        Saves some processor cycles
        """
        def present(expr):
            return (isinstance(expr.stmt, str) and
                    expr.stmt == expr.column and
                    expr.column in self.data)
        return all(present(expr) for expr in self.expressions)

    def _get_group_dataframes(self):
        """
        Get group dataframes

        Returns
        -------
        out : tuple or generator
            Group dataframes
        """
        if isinstance(self.data, GroupedDataFrame):
            grouper = self.data.groupby()
            # groupby on categorical columns uses the categories
            # even if they are not present in the data. This
            # leads to empty groups. We exclude them.
            return (gdf for _, gdf in grouper if not gdf.empty)
        else:
            return (self.data, )

    def _evaluate_expressions(self, gdfs):
        """
        Evaluate all group dataframes

        Parameters
        ----------
        gdfs : list
            Dataframes for each group

        Returns
        -------
        out : list
            Result dataframes for each group
        """
        return (self._evaluate_group_dataframe(gdf) for gdf in gdfs)

    def _evaluate_group_dataframe(self, gdf):
        """
        Evaluate a single group dataframe

        Parameters
        ----------
        gdf : pandas.DataFrame
            Input group dataframe

        Returns
        -------
        out : pandas.DataFrame
            Result data
        """
        gdf._is_copy = None
        result_index = gdf.index if self.keep_index else []
        data = pd.DataFrame(index=result_index)
        for expr in self.expressions:
            value = expr.evaluate(gdf, self.env)
            if isinstance(value, pd.DataFrame):
                data = value
                break
            else:
                _create_column(data, expr.column, value)
        data = _add_group_columns(data, gdf)
        return data

    def _concat(self, egdfs):
        """
        Concatenate evaluated group dataframes

        Parameters
        ----------
        egdfs : iterable
            Evaluated dataframes

        Returns
        -------
        edata : pandas.DataFrame
            Evaluated data
        """
        egdfs = list(egdfs)
        edata = pd.concat(egdfs, axis=0, ignore_index=False, copy=False)

        # groupby can mixup the rows. We try to maintain the original
        # order, but we can only do that if the result has a one to
        # one relationship with the original
        one2one = (
            (len(egdfs) > 1) and
            self.keep_index and
            not any(edata.index.duplicated()) and
            len(edata.index) == len(self.data.index)
        )
        if one2one:
            edata = edata.sort_index()
        else:
            edata.reset_index(drop=True, inplace=True)

        # Maybe this should happen in the verb functions
        if self.keep_groups and self.groups:
            edata = GroupedDataFrame(edata, groups=self.groups)
        return edata class Selector:
    """
    Helper to select columns of a verb
    """
    @staticmethod
    def _resolve_slices(data_columns, names):
        """
        Convert any slices into column names

        Parameters
        ----------
        data_columns : pandas.Index
            Dataframe columns
        names : tuple
            Names (including slices) of columns in the
            dataframe.

        Returns
        -------
        out : tuple
            Names of columns in the dataframe. Has no
            slices.
        """
        def _get_slice_cols(sc):
            """
            Convert slice to list of names
            """
            # Just like pandas.DataFrame.loc the stop
            # column is included
            idx_start = data_columns.get_loc(sc.start)
            idx_stop = data_columns.get_loc(sc.stop) + 1
            return data_columns[idx_start:idx_stop:sc.step]

        result = []
        for col in names:
            if isinstance(col, slice):
                result.extend(_get_slice_cols(col))
            else:
                result.append(col)
        return tuple(result)

    @classmethod
    def verify_columns(cls, selected, data_columns):
        missing_columns = selected.difference(
            data_columns,
            sort=False
        )
        if len(missing_columns):
            raise KeyError(
                "Unknown columns: "
                + ', '.join(missing_columns)
            )
        return selected

    @classmethod
    def has_minus(cls, col, data_columns):
        if isinstance(col, str) and col.startswith('-'):
            return col[1:] in data_columns
        return False

    @classmethod
    def select(cls, verb):
        """
        Return selected columns for the select verb

        Parameters
        ----------
        verb : object
            verb with the column selection attributes:

                - names
                - startswith
                - endswith
                - contains
                - matches
        """
        columns = verb.data.columns
        contains = verb.contains
        matches = verb.matches

        groups = _get_groups(verb)
        names = cls._resolve_slices(columns, verb.names)
        names_set = set(names)
        groups_set = set(groups)
        lst = [[]]

        if names and cls.has_minus(names[0], columns):
            return cls.select_minus(verb)

        if names or groups:
            # group variable missing from the selection are prepended
            missing = [g for g in groups if g not in names_set]
            missing_set = set(missing)
            c1 = missing + [x for x in names if x not in missing_set]
            lst.append(c1)

        if verb.startswith:
            c2 = [x for x in columns
                  if isinstance(x, str) and x.startswith(verb.startswith)]
            lst.append(c2)

        if verb.endswith:
            c3 = [x for x in columns if
                  isinstance(x, str) and x.endswith(verb.endswith)]
            lst.append(c3)

        if contains:
            c4 = []
            for col in columns:
                if (isinstance(col, str) and
                        any(s in col for s in contains)):
                    c4.append(col)
            lst.append(c4)

        if matches:
            c5 = []
            patterns = [x if hasattr(x, 'match') else re.compile(x)
                        for x in matches]
            for col in columns:
                if isinstance(col, str):
                    if any(bool(p.match(col)) for p in patterns):
                        c5.append(col)

            lst.append(c5)

        selected = pd.Index(list(itertools.chain(*lst))).drop_duplicates()

        if verb.drop:
            to_drop = [col for col in selected if col not in groups_set]
            selected = pd.Index(
                [col for col in columns if col not in to_drop]
            )
        return cls.verify_columns(selected, columns)

    @classmethod
    def select_minus(cls, verb):
        columns = verb.data.columns
        names = pd.Index(verb.names).drop_duplicates()
        # Columns preceeded with minus
        exclude_columns = pd.Index([
            col[1:]
            for col in names
            if cls.has_minus(col, columns)
        ])
        # Any other columns
        include_columns = pd.Index([
            col
            for col in names
            if not cls.has_minus(col, columns)
        ])
        selected = columns.difference(exclude_columns, sort=False)

        # For cases like select('-col1', 'col2', 'col1')
        # col1 has to be included
        selected = selected.append(include_columns).drop_duplicates()
        return cls.verify_columns(selected, columns)

    @classmethod
    def _all(cls, verb):
        """
        A verb
        """
        groups = set(_get_groups(verb))
        return [col for col in verb.data if col not in groups]

    @classmethod
    def _at(cls, verb):
        """
        A verb with a select text match
        """
        # Named (listed) columns are always included
        columns = cls.select(verb)
        final_columns_set = set(columns)
        groups_set = set(_get_groups(verb))
        final_columns_set -= groups_set - set(verb.names)
        return [col for col in columns if col in final_columns_set]

    @classmethod
    def _if(cls, verb):
        """
        A verb with a predicate function
        """
        pred = verb.predicate
        data = verb.data
        groups = set(_get_groups(verb))

        # force predicate
        if isinstance(pred, str):
            if not pred.endswith('_dtype'):
                pred = '{}_dtype'.format(pred)
            pred = getattr(pdtypes, pred)
        elif pdtypes.is_bool_dtype(np.array(pred)):
            # Turn boolean array into a predicate function
            it = iter(pred)

            def pred(col):
                return next(it)

        return [col for col in data
                if pred(data[col]) and col not in groups]

    @classmethod
    def get(cls, verb):
        if not hasattr(verb, 'selector') and hasattr(verb, 'names'):
            return cls.select(verb)
        elif verb.selector == '_all':
            return cls._all(verb)
        elif verb.selector == '_at':
            return cls._at(verb)
        elif verb.selector == '_if':
            return cls._if(verb) def _get_groups(verb):
    """
    Return groups
    """
    try:
        return verb.data.plydata_groups
    except AttributeError:
        return [] def _get_base_dataframe(df):
    """
    Remove all columns other than those grouped on
    """
    if isinstance(df, GroupedDataFrame):
        base_df = GroupedDataFrame(
            df.loc[:, df.plydata_groups], df.plydata_groups,
            copy=True)
    else:
        base_df = pd.DataFrame(index=df.index)
    return base_df
import warnings
import numpy as np
import pandas as pd
from ..types import GroupedDataFrame
from ..options import get_option
from ..operators import register_implementations
from ..utils import Q, get_empty_env, regular_index, unique
from .common import Evaluator, Selector
from .common import _get_groups, _get_base_dataframe
        return verb.data

    _check_modify_groups(verb)
    verb.env = verb.env.with_outer_namespace(_outer_namespace)
    with regular_index(verb.data):
        new_data = Evaluator(verb).process()
        for col in new_data:
            verb.data[col] = new_data[col]
    return verb.data


def create(verb):
    data = _get_base_dataframe(verb.data)
    verb.env = verb.env.with_outer_namespace(_outer_namespace)
    with regular_index(verb.data, data):
        new_data = Evaluator(verb, drop=True).process()
        for col in new_data:
            data[col] = new_data[col]
    return data


def sample_n(verb):
    return verb.data.sample(**verb.kwargs)


def sample_frac(verb):
    return verb.data.sample(**verb.kwargs)


def select(verb):
    columns = Selector.get(verb)
    data = verb.data.loc[:, columns]
    return data


def rename(verb):
    inplace = get_option('modify_input_data')
    data = verb.data.rename(columns=verb.lookup, inplace=inplace)
    return verb.data if inplace else data


def distinct(verb):
    data = define(verb)
    return data.drop_duplicates(subset=verb.columns,
                                keep=verb.keep)


def arrange(verb):
    # Do not evaluate if all statements correspond to
    # columns already in the dataframe
    stmts = [expr.stmt for expr in verb.expressions]
    has_all_columns = all(stmt in verb.data for stmt in stmts)
    if has_all_columns:
        df = verb.data.loc[:, stmts]
    else:
        verb.env = verb.env.with_outer_namespace({'Q': Q})
        df = Evaluator(verb, keep_index=True).process()

    if len(df.columns):
        # The index is also rearranged, but to avoid issues with
        # duplicate index values, we work with a regular index
        original_index = verb.data.index
        with regular_index(verb.data, df):
            sorted_index = df.sort_values(by=list(df.columns)).index
            data = verb.data.loc[sorted_index, :]

        if verb.reset_index:
            data.reset_index(drop=True, inplace=True)
        else:
            data.index = original_index[sorted_index]
    else:
        data = verb.data

    return data


def group_by(verb):
    verb._overwrite_groups = True
    verb.data = define(verb)

    copy = not get_option('modify_input_data')

    try:
        add = verb.add_
    except AttributeError:
        add = False

    if add:
        groups = _get_groups(verb) + verb.groups
    else:
        groups = verb.groups

    if groups:
        return GroupedDataFrame(verb.data, groups, copy=copy)
    else:
        return pd.DataFrame(verb.data, copy=copy)


def ungroup(verb):
    return pd.DataFrame(verb.data)


def group_indices(verb):
    data = verb.data
    groups = verb.groups
    if isinstance(data, GroupedDataFrame):
        if groups:
            msg = "GroupedDataFrame ignored extra groups {}"
            warnings.warn(msg.format(groups))
        else:
            groups = data.plydata_groups
    else:
        data = create(verb)
        data = GroupedDataFrame(data, groups, copy=False)

    return data.group_indices()


def summarize(verb):
    verb.env = verb.env.with_outer_namespace(_outer_namespace)

    with regular_index(verb.data):