Coverage for src/remedapy/filter.py: 68%

1 import inspect

2 from collections.abc import Callable, Iterable, Sequence

3 from typing import TypeVar, cast, overload

4

5 T = TypeVar('T')

6

7

8@overload

9 def filter(data: Iterable[T], predicate: Callable[[T], bool], /) -> Iterable[T]: ...

10

11

12@overload

13 def filter(data: Iterable[T], predicate: Callable[[T, int], bool], /) -> Iterable[T]: ...

14

15

16@overload

17 def filter(data: Sequence[T], predicate: Callable[[T, int, Sequence[T]], bool], /) -> Iterable[T]: ...

18

19

20@overload

21 def filter(predicate: Callable[[T], bool], /) -> Callable[[Iterable[T]], Iterable[T]]: ...

22

23

24@overload

25 def filter(predicate: Callable[[T, int], bool], /) -> Callable[[Iterable[T]], Iterable[T]]: ...

26

27

28@overload

29 def filter(predicate: Callable[[T, int, Sequence[T]], bool], /) -> Callable[[Sequence[T]], Iterable[T]]: ...

30

31

32 def filter(

33 data: Iterable[T] | Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool],

34 callbackfn: Callable[[T], bool] | Callable[[T, int], bool] | Callable[[T, int, Sequence[T]], bool] | None = None,

35 /,

36) -> Iterable[T] | Callable[[Iterable[T]], Iterable[T]] | Callable[[Sequence[T]], Iterable[T]]:

37 if callbackfn is None:

38 callbackfn = cast(Callable[[T], bool], data)

39

40 def inner(data: Iterable[T], /) -> Iterable[T]:

41 return filter(data, callbackfn)

42

43 return inner

44 data = cast(list[T], data)

45 sig = inspect.signature(callbackfn)

46 num_params = len(sig.parameters)

47 match num_params:

48 case 1:

49 callbackfn = cast(Callable[[T], bool], callbackfn)

50 return (item for item in data if callbackfn(item))

51 case 2:

52 callbackfn = cast(Callable[[T, int], bool], callbackfn)

53 return (item for index, item in enumerate(data) if callbackfn(item, index))

54 case 3:

55 callbackfn = cast(Callable[[T, int, Sequence[T]], bool], callbackfn)

56 return (item for index, item in enumerate(data) if callbackfn(item, index, data))

57 case _:

58 raise ValueError(f'Unsupported number of parameters: {num_params}')