diff options
Diffstat (limited to 'rumba/model.py')
| -rw-r--r-- | rumba/model.py | 21 |
1 files changed, 17 insertions, 4 deletions
diff --git a/rumba/model.py b/rumba/model.py index 68c6c51..9187fcb 100644 --- a/rumba/model.py +++ b/rumba/model.py @@ -173,6 +173,11 @@ class Node: self._validate() + def get_ipcp_by_dif(self, dif): + for ipcp in self.ipcps: + if ipcp.dif == dif: + return ipcp + def _undeclared_dif(self, dif): if dif not in self.difs: raise Exception("Invalid registration: node %s is not declared " @@ -458,9 +463,13 @@ class Experiment: for edge in dif_graphs[dif][cur]: if edge[0] not in enrolled: enrolled.add(edge[0]) + enrollee = edge[0].get_ipcp_by_dif(dif) + assert(enrollee is not None) + enroller = cur.get_ipcp_by_dif(dif) + assert(enroller is not None) self.enrollments[-1].append({'dif': dif, - 'enrollee': edge[0], - 'enroller': cur, + 'enrollee': enrollee, + 'enroller': enroller, 'lower_dif': edge[1]}) frontier.add(edge[0]) @@ -468,9 +477,13 @@ class Experiment: for cur in dif_graphs[dif]: for edge in dif_graphs[dif][cur]: if cur < edge[0]: + enrollee = cur.get_ipcp_by_dif(dif) + assert(enrollee is not None) + enroller = edge[0].get_ipcp_by_dif(dif) + assert(enroller is not None) self.enrollments[-1].append({'dif': dif, - 'enrollee': cur, - 'enroller': edge[0], + 'enrollee': enrollee, + 'enroller': enroller, 'lower_dif': edge[1]}) else: |