Coverage for smartmdao / graph.py: 100%

1from collections import defaultdict 

2from typing import List, Dict, Set 

3 

4from .models import Step 

5 

6 

7def map_producers(steps: List[Step]) -> Dict[str, Step]: 

8 """Maps each variable name to the step that produces it.""" 

9 mapping = {} 

10 for step in steps: 

11 for out in step.resolve_output_names(): 

12 mapping[out] = step 

13 return mapping 

14 

15 

16def build_dependency_graph(steps: List[Step], input_keys: Set[str], producers_map: Dict[str, Step]): 

17 """Builds a producer -> consumer adjacency list and an indegree map.""" 

18 adj_list = defaultdict(list) 

19 indegree = defaultdict(int) 

20 

21 for s in steps: 

22 indegree[s] = 0 

23 

24 for consumer in steps: 

25 # Use .get_signature() to see through decorators (e.g. @cached) 

26 sig = consumer.get_signature() 

27 for param in sig.parameters: 

28 

29 # PRIORITY: Check if it's an internal producer FIRST. 

30 if param in producers_map: 

31 producer = producers_map[param] 

32 adj_list[producer].append(consumer) 

33 indegree[consumer] += 1 

34 

35 # Only if it's NOT produced internally do we check if it's satisfied by inputs. 

36 elif param in input_keys: 

37 continue 

38 

39 return adj_list, indegree