from __future__ import annotations
import numpy as np
import typing_extensions as tp
from numpy.ma import MaskedArray
from static_frame.core.doc_str import doc_inject
from static_frame.core.exception import immutable_type_error_factory
from static_frame.core.util import (
NULL_SLICE,
TBlocKey,
TCallableAny,
TDepthLevelSpecifier,
TDtypeSpecifier,
TDtypesSpecifier,
TILocSelector,
TILocSelectorCompound,
TILocSelectorMany,
TILocSelectorOne,
TLabel,
TLocSelector,
TLocSelectorCompound,
TLocSelectorMany,
)
# from static_frame.core.util import TCallableAny
if tp.TYPE_CHECKING:
from static_frame.core.assign import Assign
from static_frame.core.batch import Batch
from static_frame.core.bus import Bus
from static_frame.core.frame import (
Frame,
FrameAssignILoc,
FrameAsType,
FrameGO,
FrameHE,
)
from static_frame.core.index import Index
from static_frame.core.index_base import IndexBase
from static_frame.core.index_hierarchy import (
IndexHierarchy,
IndexHierarchyAsType,
)
from static_frame.core.series import (
Series,
SeriesAssign,
SeriesHE,
)
from static_frame.core.type_blocks import TypeBlocks
from static_frame.core.yarn import Yarn
TNDArrayAny = np.ndarray[tp.Any, tp.Any]
TDtypeAny = np.dtype[tp.Any]
TSeriesAny = Series[tp.Any, tp.Any]
TFrameAny = Frame[tp.Any, tp.Any, tp.Unpack[tp.Tuple[tp.Any, ...]]]
TBusAny = Bus[tp.Any]
TYarnAny = Yarn[tp.Any]
# -------------------------------------------------------------------------------
TFrameOrSeries = tp.Union[
'Frame[tp.Any, tp.Any, tp.Unpack[tp.Tuple[tp.Any, ...]]]',
'Series[tp.Any, tp.Any]',
]
TVContainer_co = tp.TypeVar(
'TVContainer_co',
'Index[tp.Any]',
'Series[tp.Any, tp.Any]',
'SeriesHE[tp.Any, tp.Any]',
'Frame[tp.Any, tp.Any, tp.Unpack[tp.Tuple[tp.Any, ...]]]',
'FrameGO[tp.Any, tp.Any]',
'FrameHE[tp.Any, tp.Any, tp.Unpack[tp.Tuple[tp.Any, ...]]]',
'TypeBlocks',
'Bus[tp.Any]',
'Batch',
'Yarn[tp.Any]',
'IndexHierarchy',
'SeriesAssign',
'FrameAssignILoc',
np.ndarray[tp.Any, tp.Any],
MaskedArray[tp.Any, tp.Any],
TFrameOrSeries,
covariant=True,
)
TVIndex = tp.TypeVar('TVIndex', bound='IndexBase', default=tp.Any)
TVColumns = tp.TypeVar('TVColumns', bound='IndexBase', default=tp.Any)
TLocSelectorFunc = tp.TypeVar(
'TLocSelectorFunc',
bound=tp.Callable[[TLocSelector], TVContainer_co], # pyright: ignore
)
TILocSelectorFunc = tp.TypeVar(
'TILocSelectorFunc',
bound=tp.Callable[[TILocSelector], TVContainer_co], # pyright: ignore
)
TLocSelectorFuncInPlace = tp.TypeVar(
'TLocSelectorFuncInPlace',
bound=tp.Callable[[TLocSelector], None], # pyright: ignore
)
TILocSelectorFuncInPlace = tp.TypeVar(
'TILocSelectorFuncInPlace',
bound=tp.Callable[[TILocSelector], None], # pyright: ignore
)
TVDtype = tp.TypeVar('TVDtype', default=tp.Any)
class Interface:
__slots__ = ()
_INTERFACE: tp.Tuple[str, ...] = ()
class InterfaceBatch:
__slots__ = ()
_INTERFACE: tp.Tuple[str, ...] = ()
class InterGetItemILocReduces(Interface, tp.Generic[TVContainer_co, TVDtype]):
"""Interface for iloc selection that reduces dimensionality."""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
def __init__(
self,
func: tp.Union[
tp.Callable[[TILocSelectorOne], TVDtype],
tp.Callable[[TILocSelectorMany], TVContainer_co],
],
) -> None:
self._func: tp.Union[
tp.Callable[[TILocSelectorOne], TVDtype],
tp.Callable[[TILocSelectorMany], TVContainer_co],
] = func
@tp.overload
def __getitem__(self, key: TILocSelectorMany) -> TVContainer_co: ...
@tp.overload
def __getitem__(self, key: TILocSelectorOne) -> TVDtype: ...
def __getitem__(self, key: TILocSelector) -> tp.Any:
return self._func(key) # type: ignore
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'iloc',
key,
value,
)
class InterGetItemILoc(Interface, tp.Generic[TVContainer_co]):
"""Interface for iloc selection that does not reduce dimensionality."""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
def __init__(
self,
func: tp.Union[
tp.Callable[[TILocSelectorOne], tp.Any],
tp.Callable[[TILocSelectorMany], TVContainer_co],
],
) -> None:
self._func: tp.Union[
tp.Callable[[TILocSelectorOne], tp.Any],
tp.Callable[[TILocSelectorMany], TVContainer_co],
] = func
def __getitem__(self, key: TILocSelector) -> TVContainer_co:
return self._func(key) # type: ignore
class InterGetItemILocInPlace(Interface, tp.Generic[TVContainer_co]):
"""Variant that has None return."""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
def __init__(
self,
func: tp.Union[
tp.Callable[[TILocSelectorOne], None],
tp.Callable[[TILocSelectorMany], None],
],
) -> None:
self._func = func
def __getitem__(self, key: TILocSelector) -> None:
self._func(key) # type: ignore
class InterGetItemLocReduces(Interface, tp.Generic[TVContainer_co, TVDtype]):
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TLocSelector], TVContainer_co | TVDtype]
def __init__(
self,
func: tp.Callable[[TLocSelector], TVContainer_co | TVDtype],
) -> None:
self._func = func
@tp.overload
def __getitem__(self, key: TLocSelectorMany) -> TVContainer_co: ...
@tp.overload
def __getitem__(self, key: TLabel) -> TVDtype: ...
def __getitem__(self, key: TLocSelector) -> tp.Any:
return self._func(key)
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'loc',
key,
value,
)
class InterGetItemLoc(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TLocSelector], TVContainer_co]
def __init__(self, func: tp.Callable[[TLocSelector], TVContainer_co]) -> None:
self._func = func
def __getitem__(self, key: TLocSelector) -> TVContainer_co:
return self._func(key)
class InterGetItemLocInPlace(Interface, tp.Generic[TVContainer_co]):
"""Variant that has None return."""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TLocSelector], None]
def __init__(self, func: tp.Callable[[TLocSelector], None]) -> None:
self._func = func
def __getitem__(self, key: TLocSelector) -> None:
self._func(key)
class InterGetItemLocCompoundReduces(
Interface, tp.Generic[TVContainer_co, TVIndex, TVColumns]
):
"""Interface for compound loc selection that reduces dimensionality. TVContainer_co is the outermost container"""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TLocSelectorCompound], tp.Any]
def __init__(self, func: tp.Callable[[TLocSelectorCompound], tp.Any]) -> None:
self._func = func
@tp.overload
def __getitem__(self, key: tp.Tuple[slice, slice]) -> TVContainer_co: ...
@tp.overload # selects a Series as a row
def __getitem__(
self, key: tp.Tuple[TLabel, TLocSelectorMany]
) -> Series[TVColumns, tp.Any]: ...
@tp.overload # selects a Series as s column
def __getitem__(
self, key: tp.Tuple[TLocSelectorMany, TLabel]
) -> Series[TVIndex, tp.Any]: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[tp.List[int], tp.List[int]]
) -> TVContainer_co: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[tp.List[str], tp.List[str]]
) -> TVContainer_co: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[TLocSelectorMany, TLocSelectorMany]
) -> TVContainer_co: ...
@tp.overload
def __getitem__(self, key: tp.Tuple[TLabel, TLabel]) -> tp.Any: ...
@tp.overload
def __getitem__(self, key: TLabel) -> TSeriesAny: ...
@tp.overload
def __getitem__(self, key: TLocSelectorMany) -> TVContainer_co: ...
@tp.overload
def __getitem__(self, key: TLocSelectorCompound) -> tp.Any: ...
def __getitem__(self, key: TLocSelectorCompound) -> tp.Any:
return self._func(key)
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'loc',
key,
value,
)
class InterGetItemLocCompound(Interface, tp.Generic[TVContainer_co]):
"""Interface for compound loc selection that does not reduce dimensionality. TVContainer_co is the only delivered container container"""
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TLocSelectorCompound], TVContainer_co]
def __init__(self, func: tp.Callable[[TLocSelectorCompound], TVContainer_co]) -> None:
self._func = func
def __getitem__(self, key: TLocSelectorCompound) -> TVContainer_co:
return self._func(key)
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'loc',
key,
value,
)
class InterGetItemILocCompoundReduces(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TILocSelectorCompound], tp.Any]
def __init__(self, func: tp.Callable[[TILocSelectorCompound], tp.Any]) -> None:
self._func = func
@tp.overload
def __getitem__(self, key: TILocSelectorOne) -> TSeriesAny: ...
@tp.overload
def __getitem__(self, key: TILocSelectorMany) -> TVContainer_co: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[TILocSelectorOne, TILocSelectorMany]
) -> TSeriesAny: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[TILocSelectorMany, TILocSelectorOne]
) -> TSeriesAny: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[TILocSelectorMany, TILocSelectorMany]
) -> TVContainer_co: ...
@tp.overload
def __getitem__(
self, key: tp.Tuple[TILocSelectorOne, TILocSelectorOne]
) -> tp.Any: ...
@tp.overload
def __getitem__(self, key: TILocSelectorMany) -> TVContainer_co: ...
@tp.overload
def __getitem__(self, key: TILocSelectorCompound) -> tp.Any: ...
def __getitem__(self, key: TILocSelectorCompound) -> tp.Any:
return self._func(key)
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'iloc',
key,
value,
)
class InterGetItemILocCompound(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TILocSelectorCompound], TVContainer_co]
def __init__(
self, func: tp.Callable[[TILocSelectorCompound], TVContainer_co]
) -> None:
self._func = func
def __getitem__(self, key: TILocSelectorCompound) -> TVContainer_co:
return self._func(key)
def __setitem__(self, key: TLabel, value: tp.Any) -> None:
raise immutable_type_error_factory(
self._func.__self__.__class__, # type: ignore
'loc',
key,
value,
)
class InterfaceGetItemBLoc(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func',)
_INTERFACE = ('__getitem__',)
_func: tp.Callable[[TBlocKey], TVContainer_co]
def __init__(self, func: tp.Callable[[TBlocKey], TVContainer_co]) -> None:
self._func = func
def __getitem__(self, key: TBlocKey) -> TVContainer_co:
return self._func(key)
# -------------------------------------------------------------------------------
class InterfaceSelectDuo(Interface, tp.Generic[TVContainer_co]):
"""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: TILocSelectorFunc, func_loc: TLocSelectorFunc
) -> None:
self._func_iloc = func_iloc
self._func_loc = func_loc
@property
def iloc(self) -> InterGetItemILocReduces[TVContainer_co, tp.Any]:
return InterGetItemILocReduces(self._func_iloc)
@property
def loc(self) -> InterGetItemLocReduces[TVContainer_co, tp.Any]:
return InterGetItemLocReduces(self._func_loc) # pyright: ignore
class InterfaceSelectTrio(Interface, tp.Generic[TVContainer_co]):
"""An instance to serve as an interface to all of iloc, loc, and __getitem__ extractors. It is assumed that functionality that uses this interface returns containers that do not reduce their dimensionality."""
__slots__ = (
'_func_iloc',
'_func_loc',
'_func_getitem',
)
_INTERFACE = ('__getitem__', 'iloc', 'loc')
def __init__(
self,
*,
func_iloc: TILocSelectorFunc,
func_loc: TLocSelectorFunc,
func_getitem: TLocSelectorFunc,
) -> None:
self._func_iloc = func_iloc
self._func_loc = func_loc
self._func_getitem = func_getitem
[docs]
def __getitem__(self, key: TLocSelector) -> tp.Any:
"""Label-based selection."""
return self._func_getitem(key)
@property
def iloc(self) -> InterGetItemILoc[TVContainer_co]:
"""Integer-position based selection."""
return InterGetItemILoc(self._func_iloc)
@property
def loc(self) -> InterGetItemLoc[TVContainer_co]:
"""Label-based selection."""
return InterGetItemLoc(self._func_loc) # pyright: ignore
class InterfaceSelectQuartet(Interface, tp.Generic[TVContainer_co]):
"""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: TILocSelectorFunc,
func_loc: TLocSelectorFunc,
func_getitem: TLocSelectorFunc,
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: TLocSelector) -> tp.Any:
"""Label-based selection."""
return self._func_getitem(key)
@property
def bloc(self) -> InterGetItemLocReduces[TVContainer_co, tp.Any]:
"""Boolean based assignment."""
return InterGetItemLocReduces(self._func_bloc)
@property
def iloc(self) -> InterGetItemILocReduces[TVContainer_co, tp.Any]:
"""Integer-position based assignment."""
return InterGetItemILocReduces(self._func_iloc)
@property
def loc(self) -> InterGetItemLocReduces[TVContainer_co, tp.Any]:
"""Label-based assignment."""
return InterGetItemLocReduces(self._func_loc) # pyright: ignore
# -------------------------------------------------------------------------------
class InterfaceAssignTrio(InterfaceSelectTrio[TVContainer_co]):
"""For assignment with __getitem__, iloc, loc."""
__slots__ = ('delegate',)
def __init__(
self,
*,
func_iloc: TILocSelectorFunc,
func_loc: TLocSelectorFunc,
func_getitem: TLocSelectorFunc,
delegate: tp.Type[Assign],
) -> None:
InterfaceSelectTrio.__init__(
self,
func_iloc=func_iloc,
func_loc=func_loc,
func_getitem=func_getitem,
)
self.delegate = delegate
class InterfaceAssignQuartet(InterfaceSelectQuartet[TVContainer_co]):
"""For assignment with __getitem__, iloc, loc, bloc."""
__slots__ = ('delegate',)
def __init__(
self,
*,
func_iloc: TILocSelectorFunc,
func_loc: TLocSelectorFunc,
func_getitem: TLocSelectorFunc,
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
# -------------------------------------------------------------------------------
class InterfaceFrameAsType(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func_getitem',)
_INTERFACE = ('__getitem__', '__call__')
def __init__(self, func_getitem: tp.Callable[[TLocSelector], '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: TLocSelector) -> 'FrameAsType':
"""Selector of columns by label.
Args:
key: {key_loc}
"""
return self._func_getitem(key)
[docs]
def __call__(
self,
dtype: TDtypeSpecifier,
*,
consolidate_blocks: bool = False,
) -> TFrameAny:
"""
Apply a single ``dtype`` to all columns.
"""
return self._func_getitem(NULL_SLICE)(
dtype,
consolidate_blocks=consolidate_blocks,
)
class InterfacePersist(Interface, tp.Generic[TVContainer_co]):
__slots__ = (
'_func_iloc',
'_func_loc',
'_func_getitem',
)
_INTERFACE = ('__getitem__', 'iloc', 'loc', '__call__')
def __init__(
self,
*,
func_iloc: TILocSelectorFuncInPlace,
func_loc: TLocSelectorFuncInPlace,
func_getitem: TLocSelectorFuncInPlace,
) -> None:
self._func_iloc = func_iloc
self._func_loc = func_loc
self._func_getitem = func_getitem
[docs]
def __getitem__(self, key: TLocSelector) -> None:
"""Label-based selection."""
self._func_getitem(key)
@property
def iloc(self) -> InterGetItemILocInPlace[TVContainer_co]:
"""Integer-position based selection."""
return InterGetItemILocInPlace(self._func_iloc)
@property
def loc(self) -> InterGetItemLocInPlace[TVContainer_co]:
"""Label-based selection."""
return InterGetItemLocInPlace(self._func_loc) # pyright: ignore
[docs]
def __call__(self) -> None:
"""
Persist all `Frame`.
"""
self._func_getitem(NULL_SLICE)
class InterfaceIndexHierarchyAsType(Interface, tp.Generic[TVContainer_co]):
__slots__ = ('_func_getitem',)
_INTERFACE = ('__getitem__', '__call__')
def __init__(
self, func_getitem: tp.Callable[[TDepthLevelSpecifier], 'IndexHierarchyAsType']
) -> None:
"""
Args:
_func_getitem: a callable that expects a _func_getitem key and returns a IndexHierarchyAsType interface; for example, Frame._extract_getitem_astype.
"""
self._func_getitem = func_getitem
[docs]
@doc_inject(selector='selector')
def __getitem__(self, key: TDepthLevelSpecifier) -> 'IndexHierarchyAsType':
"""Selector of columns by label.
Args:
key: {key_loc}
"""
return self._func_getitem(key)
[docs]
def __call__(
self,
dtype: TDtypeAny,
/,
*,
consolidate_blocks: bool = False,
) -> 'IndexHierarchy':
"""
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[[TCallableAny], 'Batch'],
column_key: TLocSelector,
) -> None:
self._batch_apply = batch_apply
self._column_key = column_key
def __call__(
self,
dtypes: TDtypesSpecifier,
/,
*,
consolidate_blocks: bool = False,
) -> 'Batch':
return self._batch_apply(
lambda c: c.astype[self._column_key](
dtypes,
consolidate_blocks=consolidate_blocks,
)
)
class InterfaceBatchAsType(Interface, tp.Generic[TVContainer_co]):
"""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[[TCallableAny], 'Batch'],
) -> None:
self._batch_apply = batch_apply
[docs]
@doc_inject(selector='selector')
def __getitem__(self, key: TLocSelector) -> 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: TDtypeAny,
/,
) -> 'Batch':
"""
Apply a single ``dtype`` to all columns.
"""
return BatchAsType(
batch_apply=self._batch_apply,
column_key=NULL_SLICE,
)(dtype)
# -------------------------------------------------------------------------------
class InterfaceConsolidate(Interface, tp.Generic[TVContainer_co]):
"""An instance to serve as an interface to __getitem__ extractors."""
__slots__ = (
'_container',
'_func_getitem',
)
_INTERFACE = (
'__getitem__',
'__call__',
'status',
)
def __init__(
self,
container: TVContainer_co,
func_getitem: tp.Callable[[TLocSelector], TFrameAny],
) -> None:
"""
Args:
_func_getitem: a callable that expects a _func_getitem key and returns a Frame interface.
"""
self._container: TVContainer_co = container
self._func_getitem = func_getitem
[docs]
@doc_inject(selector='selector')
def __getitem__(self, key: TLocSelector) -> TFrameAny:
"""Return the full ``Frame``, selecting with ``key`` a subset of columns for consolidation.
Args:
key: {key_loc}
"""
return self._func_getitem(key)
[docs]
def __call__(self) -> TFrameAny:
"""
Apply consolidation to all columns.
"""
return self._func_getitem(NULL_SLICE)
@property
def status(self) -> TFrameAny:
"""Display consolidation status of this Frame."""
from static_frame.core.frame import Frame
flag_attrs: tp.Tuple[str, ...] = ('owndata', 'f_contiguous', 'c_contiguous')
columns: IndexBase = self._container.columns # type: ignore
def gen() -> tp.Iterator[tp.Sequence[tp.Any]]:
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]
iloc = iloc_start
else: # get inclusive slice
loc = slice(sub[0], sub[-1])
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(), columns=('loc', 'iloc', 'dtype', 'shape', 'ndim') + flag_attrs
)
