Source code for meshed.components

"""
Specialized components for meshed.
"""

from i2 import Sig
from typing import Callable, Any
from operator import itemgetter, attrgetter
from dataclasses import dataclass
from functools import partial




@dataclass
class Extractor:
    extractor_factory: Callable[[Any], Callable]
    extractor_params: Any
    # TODO: When migrating CI to 3.10+, can use `kw_only=True` here
    # name: str = field(default='extractor', kw_only=True)
    # input_name: str = field(default='x', kw_only=True)
    # But meanwhile, need an actual __init__ method:

    def __init__(
        self,
        extractor_factory: Callable[[Any], Callable],
        extractor_params: Any,
        *,
        name: str = 'extractor',
        input_name: str = 'x',
    ):
        self.extractor_factory = extractor_factory
        self.extractor_params = extractor_params
        self.name = name
        self.input_name = input_name
        self.__post_init__()

    def __post_init__(self):
        self.__name__ = self.name
        self.__signature__ = Sig(f'({self.input_name}, /)')
        self._call = self.extractor_factory(self.extractor_params)

    def __call__(self, x):
        return self._call(x)



def _itemgetter(items):
    if isinstance(items, str):
        items = [items]
    return itemgetter(*items)


def _attrgetter(attrs):
    if isinstance(attrs, str):
        attrs = [attrs]
    return attrgetter(*attrs)


Itemgetter = partial(Extractor, _itemgetter, name='itemgetter')
AttrGetter = partial(Extractor, _attrgetter, name='attrgetter')