@@ -137,6 +137,45 @@ static mp_obj_t set_discard(mp_obj_t self_in, mp_obj_t item) {
137137}
138138static MP_DEFINE_CONST_FUN_OBJ_2 (set_discard_obj , set_discard ) ;
139139
140+ static mp_obj_t set_diff_int (int n_args , const mp_obj_t * args , bool update ) {
141+ assert (n_args > 0 );
142+ assert (MP_OBJ_IS_TYPE (args [0 ], & set_type ));
143+ mp_obj_set_t * self ;
144+ if (update ) {
145+ self = args [0 ];
146+ } else {
147+ self = set_copy (args [0 ]);
148+ }
149+
150+
151+ for (int i = 1 ; i < n_args ; i ++ ) {
152+ mp_obj_t other = args [i ];
153+ if (self == other ) {
154+ set_clear (self );
155+ } else {
156+ mp_obj_t iter = rt_getiter (other );
157+ mp_obj_t next ;
158+ while ((next = rt_iternext (iter )) != mp_const_stop_iteration ) {
159+ set_discard (self , next );
160+ }
161+ }
162+ }
163+
164+ return self ;
165+ }
166+
167+ static mp_obj_t set_diff (int n_args , const mp_obj_t * args ) {
168+ return set_diff_int (n_args , args , false);
169+ }
170+ static MP_DEFINE_CONST_FUN_OBJ_VAR (set_diff_obj , 1 , set_diff ) ;
171+
172+ static mp_obj_t set_diff_update (int n_args , const mp_obj_t * args ) {
173+ set_diff_int (n_args , args , true);
174+ return mp_const_none ;
175+ }
176+ static MP_DEFINE_CONST_FUN_OBJ_VAR (set_diff_update_obj , 1 , set_diff_update ) ;
177+
178+
140179/******************************************************************************/
141180/* set constructors & public C API */
142181
@@ -146,6 +185,8 @@ static const mp_method_t set_type_methods[] = {
146185 { "clear" , & set_clear_obj },
147186 { "copy" , & set_copy_obj },
148187 { "discard" , & set_discard_obj },
188+ { "difference" , & set_diff_obj },
189+ { "difference_update" , & set_diff_update_obj },
149190 { NULL , NULL }, // end-of-list sentinel
150191};
151192
0 commit comments