Coverage for src/remedapy/map.py: 81%

1 import inspect

2 from collections.abc import Callable, Iterable, Sequence

3 from typing import TypeVar, cast, overload

4

5 T = TypeVar('T')

6 R = TypeVar('R')

7

8

9@overload

10 def map(data: Iterable[T], callbackfn: Callable[[T], R], /) -> Iterable[R]: ...

11

12

13@overload

14 def map(data: Iterable[T], callbackfn: Callable[[T, int], R], /) -> Iterable[R]: ...

15

16

17@overload

18 def map(data: Sequence[T], callbackfn: Callable[[T, int, Sequence[T]], R], /) -> Iterable[R]: ...

19

20

21@overload

22 def map(callbackfn: Callable[[T], R], /) -> Callable[[Iterable[T]], Iterable[R]]: ...

23

24

25@overload

26 def map(callbackfn: Callable[[T, int], R], /) -> Callable[[Iterable[T]], Iterable[R]]: ...

27

28

29@overload

30 def map(callbackfn: Callable[[T, int, Sequence[T]], R], /) -> Callable[[Sequence[T]], Iterable[R]]: ...

31

32

33 def map(

34 data: Iterable[T] | Callable[[T], R] | Callable[[T, int], R] | Callable[[T, int, Sequence[T]], R],

35 callbackfn: Callable[[T], R] | Callable[[T, int], R] | Callable[[T, int, Sequence[T]], R] | None = None,

36 /,

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

38 if callbackfn is None:

39 callbackfn = cast(Callable[[T], R], data)

40

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

42 return map(data, callbackfn)

43

44 return inner

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

46 sig = inspect.signature(callbackfn)

47 num_params = len(sig.parameters)

48 match num_params:

49 case 1:

50 callbackfn = cast(Callable[[T], R], callbackfn)

51 return (callbackfn(item) for item in data)

52 case 2:

53 callbackfn = cast(Callable[[T, int], R], callbackfn)

54 return (callbackfn(item, index) for index, item in enumerate(data))

55 case 3:

56 callbackfn = cast(Callable[[T, int, Sequence[T]], R], callbackfn)

57 return (callbackfn(item, index, data) for index, item in enumerate(data))

58 case _:

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