Coverage for smartmdao / models.py: 100%

1import inspect 

2from dataclasses import dataclass, is_dataclass 

3from typing import Callable, Dict, Optional, List, get_args, get_origin, get_type_hints 

4 

5@dataclass(eq=False) 

6class Step: 

7 """ 

8 Represents a single node in the computation graph. 

9 eq=False ensures hashability is based on object identity. 

10 """ 

11 fn: Callable 

12 manual_outputs: Optional[List[str]] = None 

13 

14 @property 

15 def name(self) -> str: 

16 return self.fn.__name__ 

17 

18 def get_signature(self) -> inspect.Signature: 

19 """ 

20 Robustly retrieves the signature of the underlying function, 

21 peeling off any decorators (like @cached) to find the real inputs. 

22 """ 

23 original_fn = inspect.unwrap(self.fn) 

24 return inspect.signature(original_fn) 

25 

26 def resolve_output_names(self) -> List[str]: 

27 """Determines variable names this step produces.""" 

28 if self.manual_outputs: 

29 return self.manual_outputs 

30 

31 # FIX: Use get_type_hints to correctly resolve string annotations  

32 # (common with 'from __future__ import annotations' or forward refs) 

33 try: 

34 original_fn = inspect.unwrap(self.fn) 

35 hints = get_type_hints(original_fn) 

36 ann = hints.get('return') 

37 except Exception: 

38 # Fallback to standard inspection if get_type_hints fails  

39 # (e.g., closures without global context) 

40 sig = self.get_signature() 

41 ann = sig.return_annotation 

42 

43 # If the function returns a Dataclass, use field names 

44 if isinstance(ann, type) and is_dataclass(ann): 

45 return list(ann.__dataclass_fields__.keys()) 

46 

47 # Default: use function name 

48 return [self.name] 

49 

50 def resolve_input_types(self) -> Dict[str, type]: 

51 """ 

52 Maps each declared parameter name to its annotated type. 

53 Parameters without a type hint are omitted (nothing to validate against). 

54 """ 

55 original_fn = inspect.unwrap(self.fn) 

56 try: 

57 hints = get_type_hints(original_fn) 

58 except Exception: 

59 # Unresolvable annotations (e.g. forward refs without global context): 

60 # skip type validation for this step rather than fail construction. 

61 return {} 

62 hints.pop('return', None) 

63 return hints 

64 

65 def resolve_output_types(self) -> Dict[str, type]: 

66 """ 

67 Maps each output variable name (see `resolve_output_names`) to its 

68 declared type, inferred from the function's return annotation: 

69 - Dataclass return -> per-field types. 

70 - Tuple return matched against manual `outputs=[...]` -> per-position types. 

71 - Single output -> the return annotation itself. 

72 Falls back to an empty mapping wherever a type can't be confidently inferred. 

73 """ 

74 original_fn = inspect.unwrap(self.fn) 

75 try: 

76 hints = get_type_hints(original_fn) 

77 except Exception: 

78 return {} 

79 

80 ann = hints.get('return') 

81 if ann is None: 

82 return {} 

83 

84 output_names = self.resolve_output_names() 

85 

86 if isinstance(ann, type) and is_dataclass(ann): 

87 field_hints = get_type_hints(ann) 

88 return {name: field_hints[name] for name in output_names if name in field_hints} 

89 

90 if self.manual_outputs and get_origin(ann) is tuple: 

91 args = get_args(ann) 

92 if len(args) == len(output_names) and Ellipsis not in args: 

93 return dict(zip(output_names, args)) 

94 return {} 

95 

96 if len(output_names) == 1: 

97 return {output_names[0]: ann} 

98 

99 return {}