Coverage for src / loman / util.py: 100%

1"""Utility functions and classes for loman computation graphs.""" 

2 

3import itertools 

4import types 

5from collections.abc import Callable, Generator, Iterable 

6from typing import Any, TypeVar 

7 

8import numpy as np 

9import pandas as pd 

10 

11T = TypeVar("T") 

12R = TypeVar("R") 

13 

14 

15def apply1( 

16 f: Callable[..., R], xs: T | list[T] | Generator[T, None, None], *args: Any, **kwds: Any 

17) -> R | list[R] | Generator[R, None, None]: 

18 """Apply function f to xs, handling generators, lists, and single values.""" 

19 if isinstance(xs, types.GeneratorType): 

20 return (f(x, *args, **kwds) for x in xs) 

21 if isinstance(xs, list): 

22 return [f(x, *args, **kwds) for x in xs] 

23 return f(xs, *args, **kwds) 

24 

25 

26def as_iterable(xs: T | Iterable[T]) -> Iterable[T]: 

27 """Convert input to iterable form if not already iterable.""" 

28 if isinstance(xs, (types.GeneratorType, list, set)): 

29 return xs # type: ignore[return-value] 

30 return (xs,) # type: ignore[return-value] 

31 

32 

33def apply_n(f: Callable[..., Any], *xs: Any, **kwds: Any) -> None: 

34 """Apply function f to the cartesian product of iterables xs.""" 

35 for p in itertools.product(*[as_iterable(x) for x in xs]): 

36 f(*p, **kwds) 

37 

38 

39class AttributeView: 

40 """Provides attribute-style access to dynamic collections.""" 

41 

42 def __init__( 

43 self, 

44 get_attribute_list: Callable[[], Iterable[str]], 

45 get_attribute: Callable[[str], Any], 

46 get_item: Callable[[Any], Any] | None = None, 

47 ) -> None: 

48 """Initialize with functions to get attribute list and individual attributes. 

49 

50 Args: 

51 get_attribute_list: Function that returns list of available attributes 

52 get_attribute: Function that takes an attribute name and returns its value 

53 get_item: Optional function for item access, defaults to get_attribute 

54 """ 

55 self.get_attribute_list = get_attribute_list 

56 self.get_attribute = get_attribute 

57 self.get_item: Callable[[Any], Any] = get_item if get_item is not None else get_attribute 

58 

59 def __dir__(self) -> list[str]: 

60 """Return list of available attributes.""" 

61 return list(self.get_attribute_list()) 

62 

63 def __getattr__(self, attr: str) -> Any: 

64 """Get attribute by name, raising AttributeError if not found.""" 

65 try: 

66 return self.get_attribute(attr) 

67 except KeyError as e: 

68 raise AttributeError(attr) from e 

69 

70 def __getitem__(self, key: Any) -> Any: 

71 """Get item by key.""" 

72 return self.get_item(key) 

73 

74 def __getstate__(self) -> dict[str, Any]: 

75 """Prepare object for serialization.""" 

76 return { 

77 "get_attribute_list": self.get_attribute_list, 

78 "get_attribute": self.get_attribute, 

79 "get_item": self.get_item, 

80 } 

81 

82 def __setstate__(self, state: dict[str, Any]) -> None: 

83 """Restore object from serialized state.""" 

84 self.get_attribute_list = state["get_attribute_list"] 

85 self.get_attribute = state["get_attribute"] 

86 self.get_item = state["get_item"] 

87 if self.get_item is None: 

88 self.get_item = self.get_attribute 

89 

90 @staticmethod 

91 def from_dict(d: dict[Any, Any], use_apply1: bool = True) -> "AttributeView": 

92 """Create an AttributeView from a dictionary.""" 

93 if use_apply1: 

94 

95 def get_attribute(xs: Any) -> Any: 

96 """Get attribute value from dictionary with apply1 support.""" 

97 return apply1(d.get, xs) 

98 else: 

99 get_attribute = d.get 

100 return AttributeView(d.keys, get_attribute) 

101 

102 

103pandas_types = (pd.Series, pd.DataFrame) 

104 

105 

106def value_eq(a: Any, b: Any) -> bool: 

107 """Compare two values for equality, handling pandas and numpy objects safely. 

108 

109 - Uses .equals for pandas Series/DataFrame 

110 - For numpy arrays, returns a single boolean using np.array_equal (treats NaNs as equal) 

111 - Falls back to == and coerces to bool when possible 

112 """ 

113 if a is b: 

114 return True 

115 

116 # pandas objects: use robust equality 

117 if isinstance(a, pandas_types): 

118 return bool(a.equals(b)) 

119 if isinstance(b, pandas_types): # pragma: no cover 

120 return bool(b.equals(a)) 

121 if isinstance(a, np.ndarray) or isinstance(b, np.ndarray): 

122 try: 

123 return bool(np.array_equal(a, b, equal_nan=True)) 

124 except Exception: 

125 return False 

126 

127 # Default comparison; ensure a single boolean 

128 try: 

129 result = a == b 

130 # If result is an array-like truth value, reduce safely 

131 if isinstance(result, (np.ndarray,)): 

132 return bool(np.all(result)) 

133 return bool(result) 

134 except Exception: 

135 return False