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

1import inspect 

2from collections.abc import Callable, Iterable, Sequence 

3from typing import TypeVar, cast, overload 

4 

5T = TypeVar('T') 

6R = TypeVar('R') 

7 

8 

9@overload 

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

11 

12 

13@overload 

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

15 

16 

17@overload 

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

19 

20 

21@overload 

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

23 

24 

25@overload 

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

27 

28 

29@overload 

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

31 

32 

33def 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}')