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

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

2 

3import itertools 

4import types 

5 

6import pandas as pd 

7 

8 

9def apply1(f, xs, *args, **kwds): 

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

11 if isinstance(xs, types.GeneratorType): 

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

13 if isinstance(xs, list): 

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

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

16 

17 

18def as_iterable(xs): 

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

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

21 return xs 

22 return (xs,) 

23 

24 

25def apply_n(f, *xs, **kwds): 

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

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

28 f(*p, **kwds) 

29 

30 

31class AttributeView: 

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

33 

34 def __init__(self, get_attribute_list, get_attribute, get_item=None): 

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

36 

37 Args: 

38 get_attribute_list: Function that returns list of available attributes 

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

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

41 """ 

42 self.get_attribute_list = get_attribute_list 

43 self.get_attribute = get_attribute 

44 self.get_item = get_item 

45 if self.get_item is None: 

46 self.get_item = get_attribute 

47 

48 def __dir__(self): 

49 """Return list of available attributes.""" 

50 return self.get_attribute_list() 

51 

52 def __getattr__(self, attr): 

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

54 try: 

55 return self.get_attribute(attr) 

56 except KeyError: 

57 raise AttributeError(attr) 

58 

59 def __getitem__(self, key): 

60 """Get item by key.""" 

61 return self.get_item(key) 

62 

63 def __getstate__(self): 

64 """Prepare object for serialization.""" 

65 return { 

66 "get_attribute_list": self.get_attribute_list, 

67 "get_attribute": self.get_attribute, 

68 "get_item": self.get_item, 

69 } 

70 

71 def __setstate__(self, state): 

72 """Restore object from serialized state.""" 

73 self.get_attribute_list = state["get_attribute_list"] 

74 self.get_attribute = state["get_attribute"] 

75 self.get_item = state["get_item"] 

76 if self.get_item is None: 

77 self.get_item = self.get_attribute 

78 

79 @staticmethod 

80 def from_dict(d, use_apply1=True): 

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

82 if use_apply1: 

83 

84 def get_attribute(xs): 

85 return apply1(d.get, xs) 

86 else: 

87 get_attribute = d.get 

88 return AttributeView(d.keys, get_attribute) 

89 

90 

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

92 

93 

94def value_eq(a, b): 

95 """Compare two values for equality, handling pandas objects specially.""" 

96 if a is b: 

97 return True 

98 if isinstance(a, pandas_types): 

99 return a.equals(b) 

100 if isinstance(b, pandas_types): 

101 return b.equals(a) 

102 try: 

103 return a == b 

104 except Exception: 

105 return False