@@ -849,19 +849,26 @@ class Grouper:
849
849
def __init__ (self , init = ()):
850
850
self ._mapping = weakref .WeakKeyDictionary (
851
851
{x : weakref .WeakSet ([x ]) for x in init })
852
+ self ._ordering = weakref .WeakKeyDictionary ()
853
+ for x in init :
854
+ if x not in self ._ordering :
855
+ self ._ordering [x ] = len (self ._ordering )
856
+ self ._next_order = len (self ._ordering ) # Plain int to simplify pickling.
852
857
853
858
def __getstate__ (self ):
854
859
return {
855
860
** vars (self ),
856
861
# Convert weak refs to strong ones.
857
862
"_mapping" : {k : set (v ) for k , v in self ._mapping .items ()},
863
+ "_ordering" : {** self ._ordering },
858
864
}
859
865
860
866
def __setstate__ (self , state ):
861
867
vars (self ).update (state )
862
868
# Convert strong refs to weak ones.
863
869
self ._mapping = weakref .WeakKeyDictionary (
864
870
{k : weakref .WeakSet (v ) for k , v in self ._mapping .items ()})
871
+ self ._ordering = weakref .WeakKeyDictionary (self ._ordering )
865
872
866
873
def __contains__ (self , item ):
867
874
return item in self ._mapping
@@ -875,10 +882,19 @@ def join(self, a, *args):
875
882
Join given arguments into the same set. Accepts one or more arguments.
876
883
"""
877
884
mapping = self ._mapping
878
- set_a = mapping .setdefault (a , weakref .WeakSet ([a ]))
879
-
885
+ try :
886
+ set_a = mapping [a ]
887
+ except KeyError :
888
+ set_a = mapping [a ] = weakref .WeakSet ([a ])
889
+ self ._ordering [a ] = self ._next_order
890
+ self ._next_order += 1
880
891
for arg in args :
881
- set_b = mapping .get (arg , weakref .WeakSet ([arg ]))
892
+ try :
893
+ set_b = mapping [arg ]
894
+ except KeyError :
895
+ set_b = mapping [arg ] = weakref .WeakSet ([arg ])
896
+ self ._ordering [arg ] = self ._next_order
897
+ self ._next_order += 1
882
898
if set_b is not set_a :
883
899
if len (set_b ) > len (set_a ):
884
900
set_a , set_b = set_b , set_a
@@ -892,9 +908,8 @@ def joined(self, a, b):
892
908
893
909
def remove (self , a ):
894
910
"""Remove *a* from the grouper, doing nothing if it is not there."""
895
- set_a = self ._mapping .pop (a , None )
896
- if set_a :
897
- set_a .remove (a )
911
+ self ._mapping .pop (a , {a }).remove (a )
912
+ self ._ordering .pop (a , None )
898
913
899
914
def __iter__ (self ):
900
915
"""
@@ -904,12 +919,12 @@ def __iter__(self):
904
919
"""
905
920
unique_groups = {id (group ): group for group in self ._mapping .values ()}
906
921
for group in unique_groups .values ():
907
- yield [ x for x in group ]
922
+ yield sorted ( group , key = self . _ordering . __getitem__ )
908
923
909
924
def get_siblings (self , a ):
910
925
"""Return all of the items joined with *a*, including itself."""
911
926
siblings = self ._mapping .get (a , [a ])
912
- return [ x for x in siblings ]
927
+ return sorted ( siblings , key = self . _ordering . get )
913
928
914
929
915
930
class GrouperView :
0 commit comments