Source code for static_frame.core.node_selector

import typing as tp

import numpy as np

from static_frame.core.assign import Assign
from static_frame.core.doc_str import doc_inject
from static_frame.core.util import NULL_SLICE
from static_frame.core.util import AnyCallable
from static_frame.core.util import DtypesSpecifier
from static_frame.core.util import GetItemKeyType

# from static_frame.core.util import AnyCallable

if tp.TYPE_CHECKING:
    from static_frame.core.batch import Batch  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.bus import Bus  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.frame import Frame  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.frame import FrameAsType  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.index import Index  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.index_hierarchy import IndexHierarchy  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.series import Series  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.series import SeriesAssign  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.type_blocks import TypeBlocks  # pylint: disable = W0611 #pragma: no cover
    from static_frame.core.yarn import Yarn  # pylint: disable = W0611 #pragma: no cover

#-------------------------------------------------------------------------------

TContainer = tp.TypeVar('TContainer',
        'Index',
        'Series',
        'Frame',
        'TypeBlocks',
        'Bus',
        'Batch',
        'Yarn',
        # 'Quilt',
        'IndexHierarchy',
        'SeriesAssign',
        )
GetItemFunc = tp.TypeVar('GetItemFunc',
        bound=tp.Callable[[GetItemKeyType], TContainer]
        )


class Interface(tp.Generic[TContainer]):
    __slots__ = ()
    INTERFACE: tp.Tuple[str, ...] = ()

class InterfaceBatch:
    __slots__ = ()
    INTERFACE: tp.Tuple[str, ...] = ()


class InterfaceGetItem(Interface[TContainer]):

    __slots__ = ('_func',)
    INTERFACE = ('__getitem__',)

    _func: tp.Callable[[GetItemKeyType], TContainer]

    def __init__(self, func: tp.Callable[[GetItemKeyType], TContainer]) -> None:
        self._func = func #type: ignore

    def __getitem__(self, key: GetItemKeyType) -> TContainer:
        return self._func(key) #type: ignore

#-------------------------------------------------------------------------------

class InterfaceSelectDuo(Interface[TContainer]):
    '''An instance to serve as an interface to all of iloc and loc
    '''

    __slots__ = (
            '_func_iloc',
            '_func_loc',
            )
    INTERFACE = ('iloc', 'loc')

    def __init__(self, *,
            func_iloc: GetItemFunc,
            func_loc: GetItemFunc) -> None:

        self._func_iloc = func_iloc
        self._func_loc = func_loc

    @property
    def iloc(self) -> InterfaceGetItem[TContainer]:
        return InterfaceGetItem(self._func_iloc)

    @property
    def loc(self) -> InterfaceGetItem[TContainer]:
        return InterfaceGetItem(self._func_loc)

class InterfaceSelectTrio(Interface[TContainer]):
    '''An instance to serve as an interface to all of iloc, loc, and __getitem__ extractors.
    '''

    __slots__ = (
            '_func_iloc',
            '_func_loc',
            '_func_getitem',
            )
    INTERFACE = ('__getitem__', 'iloc', 'loc')

    def __init__(self, *,
            func_iloc: GetItemFunc,
            func_loc: GetItemFunc,
            func_getitem: GetItemFunc,
            ) -> None:

        self._func_iloc = func_iloc
        self._func_loc = func_loc
        self._func_getitem = func_getitem

[docs] def __getitem__(self, key: GetItemKeyType) -> tp.Any: '''Label-based selection. ''' return self._func_getitem(key)
@property def iloc(self) -> InterfaceGetItem[TContainer]: '''Integer-position based selection.''' return InterfaceGetItem(self._func_iloc) @property def loc(self) -> InterfaceGetItem[TContainer]: '''Label-based selection. ''' return InterfaceGetItem(self._func_loc) class InterfaceSelectQuartet(Interface[TContainer]): '''An instance to serve as an interface to all of iloc, loc, and __getitem__ extractors. ''' __slots__ = ( '_func_iloc', '_func_loc', '_func_getitem', '_func_bloc', ) INTERFACE = ('__getitem__', 'iloc', 'loc', 'bloc') def __init__(self, *, func_iloc: GetItemFunc, func_loc: GetItemFunc, func_getitem: GetItemFunc, func_bloc: tp.Any, # not sure what is the right type ) -> None: self._func_iloc = func_iloc self._func_loc = func_loc self._func_getitem = func_getitem self._func_bloc = func_bloc def __getitem__(self, key: GetItemKeyType) -> tp.Any: '''Label-based selection. ''' return self._func_getitem(key) @property def bloc(self) -> InterfaceGetItem[TContainer]: '''Boolean based assignment.''' return InterfaceGetItem(self._func_bloc) @property def iloc(self) -> InterfaceGetItem[TContainer]: '''Integer-position based assignment.''' return InterfaceGetItem(self._func_iloc) @property def loc(self) -> InterfaceGetItem[TContainer]: '''Label-based assignment. ''' return InterfaceGetItem(self._func_loc) #------------------------------------------------------------------------------- class InterfaceAssignTrio(InterfaceSelectTrio[TContainer]): '''For assignment with __getitem__, iloc, loc. ''' __slots__ = ('delegate',) def __init__(self, *, func_iloc: GetItemFunc, func_loc: GetItemFunc, func_getitem: GetItemFunc, delegate: tp.Type[Assign] ) -> None: InterfaceSelectTrio.__init__(self, func_iloc=func_iloc, func_loc=func_loc, func_getitem=func_getitem, ) self.delegate = delegate #pylint: disable=E0237 class InterfaceAssignQuartet(InterfaceSelectQuartet[TContainer]): '''For assignment with __getitem__, iloc, loc, bloc. ''' __slots__ = ('delegate',) def __init__(self, *, func_iloc: GetItemFunc, func_loc: GetItemFunc, func_getitem: GetItemFunc, func_bloc: tp.Any, # not sure what is the right type delegate: tp.Type[Assign] ) -> None: InterfaceSelectQuartet.__init__(self, func_iloc=func_iloc, func_loc=func_loc, func_getitem=func_getitem, func_bloc=func_bloc, ) self.delegate = delegate #pylint: disable=E0237 #------------------------------------------------------------------------------- class InterfaceAsType(Interface[TContainer]): '''An instance to serve as an interface to __getitem__ extractors. Used by both :obj:`Frame` and :obj:`IndexHierarchy`. ''' __slots__ = ('_func_getitem',) INTERFACE = ('__getitem__', '__call__') def __init__(self, func_getitem: tp.Callable[[GetItemKeyType], 'FrameAsType'] ) -> None: ''' Args: _func_getitem: a callable that expects a _func_getitem key and returns a FrameAsType interface; for example, Frame._extract_getitem_astype. ''' self._func_getitem = func_getitem
[docs] @doc_inject(selector='selector') def __getitem__(self, key: GetItemKeyType) -> 'FrameAsType': '''Selector of columns by label. Args: key: {key_loc} ''' return self._func_getitem(key)
[docs] def __call__(self, dtype: np.dtype, *, consolidate_blocks: bool = False, ) -> 'Frame': ''' Apply a single ``dtype`` to all columns. ''' return self._func_getitem(NULL_SLICE)( dtype, consolidate_blocks=consolidate_blocks, )
class BatchAsType: __slots__ = ('_batch_apply', '_column_key') def __init__(self, batch_apply: tp.Callable[[AnyCallable], 'Batch'], column_key: GetItemKeyType ) -> None: self._batch_apply = batch_apply self._column_key = column_key def __call__(self, dtypes: DtypesSpecifier, *, consolidate_blocks: bool = False, ) -> 'Batch': return self._batch_apply( lambda c: c.astype[self._column_key]( dtypes, consolidate_blocks=consolidate_blocks, ) ) class InterfaceBatchAsType(Interface[TContainer]): '''An instance to serve as an interface to __getitem__ extractors. Used by both :obj:`Frame` and :obj:`IndexHierarchy`. ''' __slots__ = ('_batch_apply',) INTERFACE = ('__getitem__', '__call__') def __init__(self, batch_apply: tp.Callable[[AnyCallable], 'Batch'], ) -> None: self._batch_apply = batch_apply
[docs] @doc_inject(selector='selector') def __getitem__(self, key: GetItemKeyType) -> BatchAsType: '''Selector of columns by label. Args: key: {key_loc} ''' return BatchAsType(batch_apply=self._batch_apply, column_key=key)
[docs] def __call__(self, dtype: np.dtype) -> 'Batch': ''' Apply a single ``dtype`` to all columns. ''' return BatchAsType( batch_apply=self._batch_apply, column_key=NULL_SLICE, )(dtype)
#------------------------------------------------------------------------------- class InterfaceConsolidate(Interface[TContainer]): '''An instance to serve as an interface to __getitem__ extractors. ''' __slots__ = ( '_container', '_func_getitem', ) INTERFACE = ( '__getitem__', '__call__', 'status', ) def __init__(self, container: TContainer, func_getitem: tp.Callable[[GetItemKeyType], 'Frame'] ) -> None: ''' Args: _func_getitem: a callable that expects a _func_getitem key and returns a Frame interface. ''' self._container: TContainer = container self._func_getitem = func_getitem @doc_inject(selector='selector') def __getitem__(self, key: GetItemKeyType) -> 'Frame': '''Selector of columns by label for consolidation. Args: key: {key_loc} ''' return self._func_getitem(key) def __call__(self) -> 'Frame': ''' Apply consolidation to all columns. ''' return self._func_getitem(NULL_SLICE) @property def status(self) -> 'Frame': '''Display consolidation status of this Frame. ''' from static_frame.core.frame import Frame flag_attrs = ('owndata', 'f_contiguous', 'c_contiguous') columns = self._container.columns # type: ignore def gen() -> tp.Tuple[np.dtype, tp.Tuple[int, ...], int]: iloc_start = 0 for b in self._container._blocks._blocks: # type: ignore width = 1 if b.ndim == 1 else b.shape[1] iloc_end = iloc_start + width if iloc_end >= len(columns): iloc_slice = slice(iloc_start, None) else: iloc_slice = slice(iloc_start, iloc_end) sub = columns[iloc_slice] # returns a column iloc: tp.Union[int, slice] if len(sub) == 1: loc = sub[0] #type: ignore iloc = iloc_start else: # get inclusive slice loc = slice(sub[0], sub[-1]) #type: ignore iloc = iloc_slice yield [loc, iloc, b.dtype, b.shape, b.ndim] + [ getattr(b.flags, attr) for attr in flag_attrs] iloc_start = iloc_end return Frame.from_records(gen(),#type: ignore columns=('loc', 'iloc', 'dtype', 'shape', 'ndim') + flag_attrs )