@@ -42,7 +42,7 @@ class MigrationStep:
42
42
MigrationNodeKey = tuple [str , str ]
43
43
44
44
45
- @dataclass
45
+ @dataclass ( frozen = True )
46
46
class MigrationNode :
47
47
# TODO: @JCZuurmond the prefixes look a bit redundant
48
48
node_id : int = field (compare = False )
@@ -236,20 +236,20 @@ def generate_steps(self) -> Iterable[MigrationStep]:
236
236
# pre-compute incoming keys for best performance of self._required_step_ids
237
237
incoming = self ._invert_outgoing_to_incoming ()
238
238
incoming_counts = self ._compute_incoming_counts (incoming )
239
- key_queue = self ._create_key_queue (incoming_counts )
240
- key = key_queue .get ()
239
+ queue = self ._create_node_queue (incoming_counts )
240
+ node = queue .get ()
241
241
step_number = 1
242
- sorted_steps : list [MigrationStep ] = []
243
- while key is not None :
244
- step = self . _nodes [ key ]. as_step (step_number , sorted (n .node_id for n in incoming [key ]))
245
- sorted_steps .append (step )
242
+ ordered_steps : list [MigrationStep ] = []
243
+ while node is not None :
244
+ step = node . as_step (step_number , sorted (n .node_id for n in incoming [node . key ]))
245
+ ordered_steps .append (step )
246
246
# Update queue priorities
247
- for dependency in self ._outgoing [key ]:
247
+ for dependency in self ._outgoing [node . key ]:
248
248
incoming_counts [dependency .key ] -= 1
249
- key_queue .update (incoming_counts [dependency .key ], dependency )
249
+ queue .update (incoming_counts [dependency .key ], dependency )
250
250
step_number += 1
251
- key = key_queue .get ()
252
- return sorted_steps
251
+ node = queue .get ()
252
+ return ordered_steps
253
253
254
254
def _invert_outgoing_to_incoming (self ) -> dict [MigrationNodeKey , set [MigrationNode ]]:
255
255
result : dict [MigrationNodeKey , set [MigrationNode ]] = defaultdict (set )
@@ -266,14 +266,13 @@ def _compute_incoming_counts(
266
266
result [node_key ] = len (incoming [node_key ])
267
267
return result
268
268
269
- @staticmethod
270
- def _create_key_queue (incoming_counts : dict [tuple [str , str ], int ]) -> PriorityQueue :
271
- """Create a priority queue given the keys and their incoming counts.
269
+ def _create_node_queue (self , incoming_counts : dict [MigrationNodeKey , int ]) -> PriorityQueue :
270
+ """Create a priority queue for their nodes using the incoming count as priority.
272
271
273
272
A lower number means it is pulled from the queue first, i.e. the key with the lowest number of keys is retrieved
274
273
first.
275
274
"""
276
275
priority_queue = PriorityQueue ()
277
276
for node_key , incoming_count in incoming_counts .items ():
278
- priority_queue .put (incoming_count , node_key )
277
+ priority_queue .put (incoming_count , self . _nodes [ node_key ] )
279
278
return priority_queue
0 commit comments