Coverage for haystack/core/super_component/utils.py: 96%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from types import UnionType 

6from typing import Annotated, Any, get_args, get_origin 

7 

8from haystack.core.component.types import HAYSTACK_GREEDY_VARIADIC_ANNOTATION, HAYSTACK_VARIADIC_ANNOTATION 

9from haystack.utils.type_serialization import _build_pep604_union_type, _is_union_type 

10 

11 

12class _delegate_default: 

13 """Custom object for delegating filling of default values to the underlying components.""" 

14 

15 

16def _is_compatible( 

17 type1: type | UnionType, type2: type | UnionType, unwrap_nested: bool = True 

18) -> tuple[bool, type | UnionType | None]: 

19 """ 

20 Check if two types are compatible (bidirectional/symmetric check). 

21 

22 :param type1: First type to compare 

23 :param type2: Second type to compare 

24 :param unwrap_nested: If True, recursively unwraps nested Optional and Variadic types. 

25 If False, only unwraps at the top level. 

26 :return: Tuple of (True if types are compatible, common type if compatible) 

27 """ 

28 type1_unwrapped = _unwrap_all(type1, recursive=unwrap_nested) 

29 type2_unwrapped = _unwrap_all(type2, recursive=unwrap_nested) 

30 

31 return _types_are_compatible(type1_unwrapped, type2_unwrapped) 

32 

33 

34def _types_are_compatible(type1: type | UnionType, type2: type | UnionType) -> tuple[bool, type | UnionType | None]: 

35 """ 

36 Core type compatibility check implementing symmetric matching. 

37 

38 :param type1: First unwrapped type to compare 

39 :param type2: Second unwrapped type to compare 

40 :return: True if types are compatible, False otherwise 

41 """ 

42 # Handle Any type 

43 if type1 is Any: 

44 return True, type2 

45 if type2 is Any: 

46 return True, type1 

47 

48 # Direct equality 

49 if type1 == type2: 

50 return True, type1 

51 

52 type1_origin = get_origin(type1) 

53 type2_origin = get_origin(type2) 

54 

55 # Handle Union types (including X | Y syntax) 

56 if _is_union_type(type1_origin) or _is_union_type(type2_origin): 

57 return _check_union_compatibility(type1, type2, type1_origin, type2_origin) 

58 

59 # Handle non-Union types 

60 return _check_non_union_compatibility(type1, type2, type1_origin, type2_origin) 

61 

62 

63def _check_union_compatibility( 

64 type1: type | UnionType, type2: type | UnionType, type1_origin: Any, type2_origin: Any 

65) -> tuple[bool, type | UnionType | None]: 

66 """Handle all Union type compatibility cases (including X | Y syntax).""" 

67 if _is_union_type(type1_origin) and not _is_union_type(type2_origin): 

68 # Find all compatible types from the union 

69 compatible_types = [] 

70 for union_arg in get_args(type1): 

71 is_compat, common = _types_are_compatible(union_arg, type2) 

72 if is_compat and common is not None: 

73 compatible_types.append(common) 

74 if compatible_types: 

75 result_type = _build_pep604_union_type(compatible_types) 

76 return True, result_type 

77 return False, None 

78 

79 if _is_union_type(type2_origin) and not _is_union_type(type1_origin): 

80 # Find all compatible types from the union 

81 compatible_types = [] 

82 for union_arg in get_args(type2): 

83 is_compat, common = _types_are_compatible(type1, union_arg) 

84 if is_compat and common is not None: 

85 compatible_types.append(common) 

86 if compatible_types: 

87 result_type = _build_pep604_union_type(compatible_types) 

88 return True, result_type 

89 return False, None 

90 

91 # Both are Union types 

92 compatible_types = [] 

93 for arg1 in get_args(type1): 

94 for arg2 in get_args(type2): 

95 is_compat, common = _types_are_compatible(arg1, arg2) 

96 if is_compat and common is not None: 

97 compatible_types.append(common) 

98 

99 if compatible_types: 

100 result_type = _build_pep604_union_type(compatible_types) 

101 return True, result_type 

102 return False, None 

103 

104 

105def _check_non_union_compatibility( 

106 type1: type | UnionType, type2: type | UnionType, type1_origin: Any, type2_origin: Any 

107) -> tuple[bool, type | UnionType | None]: 

108 """Handle non-Union type compatibility cases.""" 

109 # If no origin, compare types directly 

110 if not type1_origin and not type2_origin: 

111 if type1 == type2: 

112 return True, type1 

113 return False, None 

114 

115 # Both must have origins and they must be equal 

116 if not (type1_origin and type2_origin and type1_origin == type2_origin): 

117 return False, None 

118 

119 # Compare generic type arguments 

120 type1_args = get_args(type1) 

121 type2_args = get_args(type2) 

122 

123 if len(type1_args) != len(type2_args): 

124 return False, None 

125 

126 # Check if all arguments are compatible 

127 common_args = [] 

128 for t1_arg, t2_arg in zip(type1_args, type2_args, strict=True): 

129 is_compat, common = _types_are_compatible(t1_arg, t2_arg) 

130 if not is_compat: 

131 return False, None 

132 common_args.append(common) 

133 

134 # Reconstruct the type with common arguments 

135 return True, type1_origin[tuple(common_args)] 

136 

137 

138def _unwrap_all(t: type | UnionType, recursive: bool) -> type | UnionType: 

139 """ 

140 Unwrap a type until no more unwrapping is possible. 

141 

142 :param t: Type to unwrap 

143 :param recursive: If True, recursively unwraps nested types 

144 :return: The fully unwrapped type 

145 """ 

146 # First handle top-level Variadic/GreedyVariadic 

147 if _is_variadic_type(t): 

148 t = _unwrap_variadics(t, recursive=recursive) 

149 else: 

150 # If it's a generic type and we're unwrapping recursively 

151 origin = get_origin(t) 

152 if recursive and origin is not None and (args := get_args(t)): 

153 unwrapped_args = tuple(_unwrap_all(arg, recursive) for arg in args) 

154 # types.UnionType (PEP 604 X | Y) is not subscriptable, so we use _build_pep604_union_type 

155 if origin is UnionType: 

156 t = _build_pep604_union_type(list(unwrapped_args)) 

157 else: 

158 t = origin[unwrapped_args] 

159 

160 return t 

161 

162 

163def _is_variadic_type(t: type | UnionType) -> bool: 

164 """Check if type is a Variadic or GreedyVariadic type.""" 

165 origin = get_origin(t) 

166 if origin is Annotated: 

167 args = get_args(t) 

168 return len(args) >= 2 and args[1] in (HAYSTACK_VARIADIC_ANNOTATION, HAYSTACK_GREEDY_VARIADIC_ANNOTATION) # noqa: PLR2004 

169 return False 

170 

171 

172def _unwrap_variadics(t: type | UnionType, recursive: bool) -> type | UnionType: 

173 """ 

174 Unwrap Variadic or GreedyVariadic annotated types. 

175 

176 :param t: Type to unwrap 

177 :param recursive: If True, recursively unwraps nested types 

178 :return: Unwrapped type if it was a variadic type, original type otherwise 

179 """ 

180 if not _is_variadic_type(t): 

181 return t 

182 

183 args = get_args(t) 

184 # Get the Iterable[X] type and extract X 

185 iterable_type = args[0] 

186 inner_type = get_args(iterable_type)[0] 

187 

188 # Only recursively unwrap if requested 

189 if recursive: 

190 return _unwrap_all(inner_type, recursive) 

191 return inner_type