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 Expression(*args, **kwargs):
    """
    Return an appropriate Expression given the arguments

    Parameters
    ----------
    args : tuple
        Positional arguments passed to the Expression class
    kwargs : dict
        Keyword arguments passed to the Expression class
    """
    # dispatch
    if not hasattr(args[0], '_Expression'):
        return BaseExpression(*args, *kwargs)
    else:
        return args[0]._Expression(*args, **kwargs) 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 _get_groups(verb):
    """
    Return groups
    """
    try:
        return verb.data.plydata_groups
    except AttributeError:
        return [] 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 build_expressions(verb):
    """
    Build expressions for helper verbs

    Parameters
    ----------
    verb : verb
        A verb with a *functions* attribute.

    Returns
    -------
    out : tuple
        (List of Expressions, New columns). The expressions and the
        new columns in which the results of those expressions will
        be stored. Even when a result will stored in a column with
        an existing label, that column is still considered new,
        i.e An expression ``x='x+1'``, will create a new_column `x`
        to replace an old column `x`.
    """
    def partial(func, col, *args, **kwargs):
        """
        Make a function that acts on a column in a dataframe

        Parameters
        ----------
        func : callable
            Function
        col : str
            Column
        args : tuple
            Arguments to pass to func
        kwargs : dict
            Keyword arguments to func

        Results
        -------
        new_func : callable
            Function that takes a dataframe, and calls the
            original function on a column in the dataframe.
        """
        def new_func(gdf):
            return func(gdf[col], *args, **kwargs)

        return new_func

    def make_statement(func, col):
        """
        A statement of function called on a column in a dataframe

        Parameters
        ----------
        func : str or callable
            Function to call on a dataframe column
        col : str
            Column
        """
        if isinstance(func, str):
            expr = '{}({})'.format(func, col)
        elif callable(func):
            expr = partial(func, col, *verb.args, **verb.kwargs)
        else:
            raise TypeError("{} is not a function".format(func))
        return expr

    def func_name(func):
        """
        Return name of a function.

        If the function is `np.sin`, we return `sin`.
        """
        if isinstance(func, str):
            return func

        try:
            return func.__name__
        except AttributeError:
            return ''

    # Generate function names. They act as identifiers (postfixed
    # to the original columns) in the new_column names.
    if isinstance(verb.functions, (tuple, list)):
        names = (func_name(func) for func in verb.functions)
        names_and_functions = zip(names, verb.functions)
    else:
        names_and_functions = verb.functions.items()

    # Create statements for the expressions
    # and postfix identifiers
    columns = Selector.get(verb)  # columns to act on
    postfixes = []
    stmts = []
    for name, func in names_and_functions:
        postfixes.append(name)
        for col in columns:
            stmts.append(make_statement(func, col))

    if not stmts:
        stmts = columns

    # Names of the new columns
    # e.g col1_mean, col2_mean, col1_std, col2_std
    add_postfix = (isinstance(verb.functions, dict) or
                   len(verb.functions) > 1)
    if add_postfix:
        fmt = '{}_{}'.format
        new_columns = [fmt(c, p) for p in postfixes for c in columns]
    else:
        new_columns = columns

    expressions = [Expression(stmt, col)
                   for stmt, col in zip(stmts, new_columns)]
    return expressions, new_columns 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 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 define(verb):
    if not get_option('modify_input_data'):
        verb.data = verb.data.copy()

    if not verb.expressions:
        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 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 _check_modify_groups(verb):
def define(verb):
def create(verb):
def sample_n(verb):
def sample_frac(verb):
def select(verb):
def rename(verb):
def distinct(verb):
def arrange(verb):
def group_by(verb):
def ungroup(verb):
def group_indices(verb):
def summarize(verb):
def query(verb):
def do(verb):
def head(verb):
def tail(verb):
def pull(verb):
def slice_rows(verb):
def _nth(arr, n):
def _n_distinct(arr):
from functools import wraps
from ..types import GroupedDataFrame
from ..expressions import Expression
from ..operators import register_implementations
from .common import _get_groups, Selector, build_expressions
from .one_table import arrange, create, define, group_by
from .one_table import mutate, rename, summarize
import pandas as pd
import numpy as np
        return verb.func(verb.data, *verb.args, **verb.kwargs)


def tally(verb):
    # Prepare for summarize
    if verb.weights is not None:
        if isinstance(verb.weights, str):
            # Summarize will evaluate and sum up the weights
            stmt = 'sum({})'.format(verb.weights)
        else:
            # Do the summation here. The result does not depend
            # on the dataframe.
            stmt = np.sum(verb.weights)
    else:
        stmt = 'n()'

    verb.expressions = [Expression(stmt, 'n')]
    data = summarize(verb)
    if verb.sort:
        data = data.sort_values(by='n', ascending=False)
        data.reset_index(drop=True, inplace=True)

    return data


def count(verb):
    groups = _get_groups(verb)

    # When grouping, add to the any current groups
    verb.add_ = True
    verb.data = group_by(verb)
    data = tally(verb)

    # Restore original groups
    if groups:
        data = GroupedDataFrame(data, groups, copy=False)
    return data


def add_tally(verb):
    if verb.weights:
        if isinstance(verb.weights, str):
            stmt = 'sum({})'.format(verb.weights)
        else:
            stmt = verb.weights
    else:
        stmt = 'n()'

    verb.expressions = [Expression(stmt, 'n')]
    data = define(verb)
    if verb.sort:
        data = data.sort_values(by='n')
        data.reset_index(drop=True, inplace=True)

    return data


def add_count(verb):
    groups = _get_groups(verb)

    # When grouping, add to the any current groups
    verb.add_ = True
    verb.data = group_by(verb)
    data = add_tally(verb)

    if groups:
        # Restore original groups
        data = GroupedDataFrame(data, groups, copy=False)
    else:
        # Remove counted groups
        data = pd.DataFrame(data, copy=False)
    return data


def _query_helper(verb):
    columns = Selector.get(verb)

    # Wrap the predicate in brackets
    fmt = '({})'.format(verb.vars_predicate).format

    # Create tokens expressions for each column
    tokens = [fmt(_=col) for col in columns]

    # Join the tokens into a compound expression
    sep = ' | ' if verb.any_vars else ' & '
    compound_expr = sep.join(tokens)

    # Create final expression, execute and get boolean selectors
    # for the rows
    result_col = '_plydata_dummy_col_'
    verb.expressions = [Expression(compound_expr, result_col)]
    _data = create(verb)
    bool_idx = _data[result_col].values

    data = verb.data.loc[bool_idx, :]
    if verb.reset_index:
        data.reset_index(drop=True, inplace=True)
    return data


def _rename_helper(verb):
    # The selector does all the work for the _all, _at & _if
    columns = Selector.get(verb)
    groups = set(_get_groups(verb))
    columns = [c for c in columns if c not in groups]
    func = verb.functions[0]
    new_names = [func(name) for name in columns]
    verb.lookup = {old: new for old, new in zip(columns, new_names)}
    return rename(verb)


def _select_helper(verb):
    # The selector does all the work for the _all, _at & _if
    columns = Selector.get(verb)
    groups = _get_groups(verb)
    groups_set = set(groups)
    columns = [c for c in columns if c not in groups_set]
    func = verb.functions[0]
    new_names = [func(name) for name in columns]
    verb.lookup = {old: new for old, new in zip(columns, new_names)}

    data = rename(verb)