Coverage for core / src / sensorkit / common / graph.py: 98%

1# SPDX-License-Identifier: Apache-2.0 

2"""Graph utilities for dependency-aware state election and topological evaluation.""" 

3 

4from __future__ import annotations 

5 

6import collections 

7from collections.abc import Hashable, Mapping, Sequence 

8from dataclasses import dataclass 

9 

10type AdjacencyMap[E: Hashable] = Mapping[E, Sequence[E]] 

11 

12 

13class StateElection[E: Hashable = str]: 

14 """Dependency-aware up/down state election.""" 

15 

16 @dataclass 

17 class VoteTally: 

18 """Accumulated up and down vote counts for a single election subject.""" 

19 

20 upvotes: int = 0 

21 downvotes: int = 0 

22 

23 def __init__(self, dependency_graph: AdjacencyMap[E]): 

24 self._graph = dependency_graph 

25 self._votes: dict[str | None, dict[E, bool | None]] = collections.defaultdict(dict) 

26 self._tally: dict[E, StateElection.VoteTally] = collections.defaultdict(self.VoteTally) 

27 

28 def vote(self, source: str | None = None, *, subject: E, vote: bool | None): 

29 """Set or clear an up/down state vote for the specified subject. 

30 

31 The `source` parameter specifies a "voter" for this vote. When multiple votes for the same 

32 subject exist, pessimistic semantics apply: a single downvote overrules any number of 

33 upvotes. 

34 """ 

35 votes = self._votes[source] 

36 prev_vote = votes.get(subject, None) 

37 

38 if vote is prev_vote: 

39 return 

40 

41 tally = self._tally[subject] 

42 

43 match prev_vote: 

44 case True: 

45 tally.upvotes -= 1 

46 case False: 

47 tally.downvotes -= 1 

48 

49 votes[subject] = vote 

50 

51 match vote: 

52 case True: 

53 tally.upvotes += 1 

54 case False: 

55 tally.downvotes += 1 

56 

57 def _evaluate_subject(self, subject: E, output: dict[E, bool], *, dependency: bool = False): 

58 if subject in output: 

59 return output[subject] 

60 

61 tally = self._tally[subject] 

62 

63 if tally.downvotes > 0: 

64 # There was at least one direct downvote for this subject, which takes precedence. 

65 output[subject] = False 

66 elif tally.upvotes + dependency > 0: 

67 # There is at least one upvote and/or this subject is a dependency of another. Now 

68 # check our own dependencies. 

69 deps_good = True 

70 

71 if subject in self._graph: 

72 for other in self._graph[subject]: 

73 # Recurse to evaluating the dependency subject. 

74 deps_good &= self._evaluate_subject(other, output, dependency=True) 

75 

76 output[subject] = deps_good 

77 else: 

78 return None 

79 

80 return output[subject] 

81 

82 def evaluate(self): 

83 """Return the state election results. 

84 

85 Returns a dictionary mapping subjects to their evaluated boolean states. The keys of the 

86 dictionary are in the topological order of the subject dependency graph. Subjects with no 

87 votes will be absent from the dictionary. 

88 """ 

89 output: dict[E, bool] = {} 

90 

91 # Topological sort to evaluate each element's consensus vote with dependencies. 

92 for subject in self._graph: 

93 self._evaluate_subject(subject, output) 

94 

95 return output