<commit_before>import collections

import numpy as np

from pandas import Panel, DataFrame, MultiIndex, Series, Timestamp
from pandas.core.indexing import _NDFrameIndexer

from trtools.monkey import patch, patch_prop

KEY_TRANS = {}

@patch([DataFrame], '__getitem__', override=True)
def df__getitem__(self, key):
    """
        This is different because we are supporting objects as keys. 
        Some keys might match both a basestring and int. The `in` keyword
        will match by hash so can only match one type.
    """
    try:
        key = KEY_TRANS.get(key, key)
        col = self._old___getitem__(key) 
        return col
    except:
        pass

    return self._old___getitem__(key) 

class TransFrameIndexer(object):
    """
        Taking over .ix
    """
    def __init__(self, obj):
        self.obj = obj

    def __getitem__(self, key):
        ix = self.obj._ix
        if type(key) is tuple:
            try:
                return self.obj.get_value(*key)
            except Exception:
                pass

            if len(key) >= 2:
                key = process_cols(self.obj, key)
            return ix._getitem_tuple(key)
        else:
            return ix._getitem_axis(key, axis=0)

    def __setitem__(self, key, value):
        self.obj._ix[key] = value

    def __getattr__(self, key):
        return getattr(self.obj._ix, key)

def process_cols(obj, key):
    """
    """

    # don't handle slice columns
    if isinstance(key[1], slice):
        return key

    cols = []
    columns = key[1]

    # single key
    if not isinstance(columns, collections.Iterable) \
       or isinstance(columns, basestring):
        ind = obj.columns.get_loc(columns)
        cols = obj.columns[ind]
    else:
        columns = [KEY_TRANS.get(c, c) for c in columns]
        mask = obj.columns.isin(columns).nonzero()[0]
        cols = obj.columns[mask]

    new_key = [key[0]]
    new_key.append(cols)
    new_key.extend(key[2:])
    return tuple(new_key)

@patch_prop([DataFrame], 'ix')
def ix(self):
    """
    """
    if not hasattr(self, '_ix') or self._ix is None:
        self._ix = _NDFrameIndexer(self)<commit_msg>added exception for string key<commit_after>import collections

import numpy as np

from pandas import Panel, DataFrame, MultiIndex, Series, Timestamp
from pandas.core.indexing import _NDFrameIndexer

from trtools.monkey import patch, patch_prop

KEY_TRANS = {}

@patch([DataFrame], '__getitem__', override=True)
def df__getitem__(self, key):
    """
        This is different because we are supporting objects as keys. 
        Some keys might match both a basestring and int. The `in` keyword
        will match by hash so can only match one type.
    """
    # Logic is to keep supporting strings even if they have an 
    # entry in translate table
    # Not sure if supporting ints is worthwhile since
    # int columns can exist automatically
    if isinstance(key, basestring) and key in self.columns:
        return self._old___getitem__(key) 

    try:
        key = KEY_TRANS.get(key, key)
        col = self._old___getitem__(key) 
        return col
    except:
        pass

    return self._old___getitem__(key) 

class TransFrameIndexer(object):
    """
        Taking over .ix
    """
    def __init__(self, obj):
        self.obj = obj

    def __getitem__(self, key):
        ix = self.obj._ix
        if type(key) is tuple:
            try:
                return self.obj.get_value(*key)
            except Exception:
                pass

            if len(key) >= 2:
                key = process_cols(self.obj, key)
            return ix._getitem_tuple(key)
        else:
            return ix._getitem_axis(key, axis=0)

    def __setitem__(self, key, value):
        self.obj._ix[key] = value

    def __getattr__(self, key):
        return getattr(self.obj._ix, key)

def process_cols(obj, key):
    """
    """

    # don't handle slice columns
    if isinstance(key[1], slice):
        return key

    cols = []
    columns = key[1]

    # single key
    if not isinstance(columns, collections.Iterable) \
       or isinstance(columns, basestring):
        ind = obj.columns.get_loc(columns)
        cols = obj.columns[ind]
    else:
        columns = [KEY_TRANS.get(c, c) for c in columns]
        cols = [obj.columns[obj.columns.get_loc(c)] for c in columns]

    new_key = [key[0]]
    new_key.append(cols)
    new_key.extend(key[2:])
    return tuple(new_key)

@patch_prop([DataFrame], 'ix')
def ix(self):
    """
    """
    if not hasattr(self, '_ix') or self._ix is None:
        self._ix = _NDFrameIndexer(self)
{'fraction_non_alphanumeric': 0.10306588388780169, 'fraction_numerical': 0.002826701456838443, 'mean_word_length': 2.311735061195104, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3724546', 'n_tokens_mistral': 1453, 'n_tokens_neox': 1399, 'n_words': 408}