@@ -50,13 +50,41 @@ typedef struct _mp_obj_set_it_t {
5050
5151STATIC mp_obj_t set_it_iternext (mp_obj_t self_in );
5252
53+ STATIC bool is_set_or_frozenset (mp_obj_t o ) {
54+ return MP_OBJ_IS_TYPE (o , & mp_type_set ) || MP_OBJ_IS_TYPE (o , & mp_type_frozenset );
55+ }
56+
57+ STATIC void check_set_or_frozenset (mp_obj_t o ) {
58+ if (!is_set_or_frozenset (o )) {
59+ nlr_raise (mp_obj_new_exception_msg_varg (& mp_type_TypeError , "'set' object required" ));
60+ }
61+ }
62+
63+ STATIC void check_set (mp_obj_t o ) {
64+ if (!MP_OBJ_IS_TYPE (o , & mp_type_set )) {
65+ // Emulate CPython behavior
66+ // AttributeError: 'frozenset' object has no attribute 'add'
67+ if (MP_OBJ_IS_TYPE (o , & mp_type_frozenset )) {
68+ nlr_raise (mp_obj_new_exception_msg (& mp_type_AttributeError , "'frozenset' has no such attribute" ));
69+ }
70+ nlr_raise (mp_obj_new_exception_msg_varg (& mp_type_TypeError , "'set' object required" ));
71+ }
72+ }
73+
5374STATIC void set_print (void (* print )(void * env , const char * fmt , ...), void * env , mp_obj_t self_in , mp_print_kind_t kind ) {
5475 mp_obj_set_t * self = self_in ;
76+ bool is_frozen = MP_OBJ_IS_TYPE (self_in , & mp_type_frozenset );
5577 if (self -> set .used == 0 ) {
78+ if (is_frozen ) {
79+ print (env , "frozen" );
80+ }
5681 print (env , "set()" );
5782 return ;
5883 }
5984 bool first = true;
85+ if (is_frozen ) {
86+ print (env , "frozenset(" );
87+ }
6088 print (env , "{" );
6189 for (int i = 0 ; i < self -> set .alloc ; i ++ ) {
6290 if (MP_SET_SLOT_IS_FILLED (& self -> set , i )) {
@@ -68,6 +96,9 @@ STATIC void set_print(void (*print)(void *env, const char *fmt, ...), void *env,
6896 }
6997 }
7098 print (env , "}" );
99+ if (is_frozen ) {
100+ print (env , ")" );
101+ }
71102}
72103
73104
@@ -82,12 +113,14 @@ STATIC mp_obj_t set_make_new(mp_obj_t type_in, uint n_args, uint n_kw, const mp_
82113 case 1 :
83114 {
84115 // 1 argument, an iterable from which we make a new set
85- mp_obj_t set = mp_obj_new_set (0 , NULL );
116+ mp_obj_set_t * set = mp_obj_new_set (0 , NULL );
86117 mp_obj_t iterable = mp_getiter (args [0 ]);
87118 mp_obj_t item ;
88119 while ((item = mp_iternext (iterable )) != MP_OBJ_STOP_ITERATION ) {
89120 mp_obj_set_store (set , item );
90121 }
122+ // Set actual set/frozenset type
123+ set -> base .type = type_in ;
91124 return set ;
92125 }
93126
@@ -132,15 +165,15 @@ STATIC mp_obj_t set_getiter(mp_obj_t set_in) {
132165/* set methods */
133166
134167STATIC mp_obj_t set_add (mp_obj_t self_in , mp_obj_t item ) {
135- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
168+ check_set ( self_in );
136169 mp_obj_set_t * self = self_in ;
137170 mp_set_lookup (& self -> set , item , MP_MAP_LOOKUP_ADD_IF_NOT_FOUND );
138171 return mp_const_none ;
139172}
140173STATIC MP_DEFINE_CONST_FUN_OBJ_2 (set_add_obj , set_add );
141174
142175STATIC mp_obj_t set_clear (mp_obj_t self_in ) {
143- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
176+ check_set ( self_in );
144177 mp_obj_set_t * self = self_in ;
145178
146179 mp_set_clear (& self -> set );
@@ -149,8 +182,7 @@ STATIC mp_obj_t set_clear(mp_obj_t self_in) {
149182}
150183STATIC MP_DEFINE_CONST_FUN_OBJ_1 (set_clear_obj , set_clear );
151184
152- STATIC mp_obj_t set_copy (mp_obj_t self_in ) {
153- assert (MP_OBJ_IS_TYPE (self_in , & mp_type_set ));
185+ STATIC mp_obj_t set_copy_as_mutable (mp_obj_t self_in ) {
154186 mp_obj_set_t * self = self_in ;
155187
156188 mp_obj_set_t * other = m_new_obj (mp_obj_set_t );
@@ -161,10 +193,20 @@ STATIC mp_obj_t set_copy(mp_obj_t self_in) {
161193
162194 return other ;
163195}
196+
197+ STATIC mp_obj_t set_copy (mp_obj_t self_in ) {
198+ check_set_or_frozenset (self_in );
199+ mp_obj_set_t * self = self_in ;
200+
201+ mp_obj_set_t * other = set_copy_as_mutable (self );
202+ other -> base .type = self -> base .type ;
203+
204+ return other ;
205+ }
164206STATIC MP_DEFINE_CONST_FUN_OBJ_1 (set_copy_obj , set_copy );
165207
166208STATIC mp_obj_t set_discard (mp_obj_t self_in , mp_obj_t item ) {
167- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
209+ check_set ( self_in );
168210 mp_obj_set_t * self = self_in ;
169211 mp_set_lookup (& self -> set , item , MP_MAP_LOOKUP_REMOVE_IF_FOUND );
170212 return mp_const_none ;
@@ -173,12 +215,14 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_discard_obj, set_discard);
173215
174216STATIC mp_obj_t set_diff_int (int n_args , const mp_obj_t * args , bool update ) {
175217 assert (n_args > 0 );
176- assert ( MP_OBJ_IS_TYPE ( args [ 0 ], & mp_type_set ));
218+
177219 mp_obj_set_t * self ;
178220 if (update ) {
221+ check_set (args [0 ]);
179222 self = args [0 ];
180223 } else {
181- self = set_copy (args [0 ]);
224+ check_set_or_frozenset (args [0 ]);
225+ self = set_copy_as_mutable (args [0 ]);
182226 }
183227
184228
@@ -195,6 +239,7 @@ STATIC mp_obj_t set_diff_int(int n_args, const mp_obj_t *args, bool update) {
195239 }
196240 }
197241
242+ self -> base .type = ((mp_obj_set_t * )args [0 ])-> base .type ;
198243 return self ;
199244}
200245
@@ -210,7 +255,12 @@ STATIC mp_obj_t set_diff_update(uint n_args, const mp_obj_t *args) {
210255STATIC MP_DEFINE_CONST_FUN_OBJ_VAR (set_diff_update_obj , 1 , set_diff_update );
211256
212257STATIC mp_obj_t set_intersect_int (mp_obj_t self_in , mp_obj_t other , bool update ) {
213- assert (MP_OBJ_IS_TYPE (self_in , & mp_type_set ));
258+ if (update ) {
259+ check_set (self_in );
260+ } else {
261+ check_set_or_frozenset (self_in );
262+ }
263+
214264 if (self_in == other ) {
215265 return update ? mp_const_none : set_copy (self_in );
216266 }
@@ -247,7 +297,7 @@ STATIC mp_obj_t set_intersect_update(mp_obj_t self_in, mp_obj_t other) {
247297STATIC MP_DEFINE_CONST_FUN_OBJ_2 (set_intersect_update_obj , set_intersect_update );
248298
249299STATIC mp_obj_t set_isdisjoint (mp_obj_t self_in , mp_obj_t other ) {
250- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
300+ check_set_or_frozenset ( self_in );
251301 mp_obj_set_t * self = self_in ;
252302
253303 mp_obj_t iter = mp_getiter (other );
@@ -264,7 +314,7 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_2(set_isdisjoint_obj, set_isdisjoint);
264314STATIC mp_obj_t set_issubset_internal (mp_obj_t self_in , mp_obj_t other_in , bool proper ) {
265315 mp_obj_set_t * self ;
266316 bool cleanup_self = false;
267- if (MP_OBJ_IS_TYPE (self_in , & mp_type_set )) {
317+ if (is_set_or_frozenset (self_in )) {
268318 self = self_in ;
269319 } else {
270320 self = set_make_new ((mp_obj_t )& mp_type_set , 1 , 0 , & self_in );
@@ -273,7 +323,7 @@ STATIC mp_obj_t set_issubset_internal(mp_obj_t self_in, mp_obj_t other_in, bool
273323
274324 mp_obj_set_t * other ;
275325 bool cleanup_other = false;
276- if (MP_OBJ_IS_TYPE (other_in , & mp_type_set )) {
326+ if (is_set_or_frozenset (other_in )) {
277327 other = other_in ;
278328 } else {
279329 other = set_make_new ((mp_obj_t )& mp_type_set , 1 , 0 , & other_in );
@@ -292,6 +342,7 @@ STATIC mp_obj_t set_issubset_internal(mp_obj_t self_in, mp_obj_t other_in, bool
292342 }
293343 }
294344 }
345+ // TODO: Should free objects altogether
295346 if (cleanup_self ) {
296347 set_clear (self );
297348 }
@@ -319,9 +370,9 @@ STATIC mp_obj_t set_issuperset_proper(mp_obj_t self_in, mp_obj_t other_in) {
319370}
320371
321372STATIC mp_obj_t set_equal (mp_obj_t self_in , mp_obj_t other_in ) {
322- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
373+ check_set_or_frozenset ( self_in );
323374 mp_obj_set_t * self = self_in ;
324- if (!MP_OBJ_IS_TYPE (other_in , & mp_type_set )) {
375+ if (!is_set_or_frozenset (other_in )) {
325376 return mp_const_false ;
326377 }
327378 mp_obj_set_t * other = other_in ;
@@ -332,7 +383,7 @@ STATIC mp_obj_t set_equal(mp_obj_t self_in, mp_obj_t other_in) {
332383}
333384
334385STATIC mp_obj_t set_pop (mp_obj_t self_in ) {
335- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
386+ check_set ( self_in );
336387 mp_obj_set_t * self = self_in ;
337388 mp_obj_t obj = mp_set_remove_first (& self -> set );
338389 if (obj == MP_OBJ_NULL ) {
@@ -343,7 +394,7 @@ STATIC mp_obj_t set_pop(mp_obj_t self_in) {
343394STATIC MP_DEFINE_CONST_FUN_OBJ_1 (set_pop_obj , set_pop );
344395
345396STATIC mp_obj_t set_remove (mp_obj_t self_in , mp_obj_t item ) {
346- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
397+ check_set ( self_in );
347398 mp_obj_set_t * self = self_in ;
348399 if (mp_set_lookup (& self -> set , item , MP_MAP_LOOKUP_REMOVE_IF_FOUND ) == MP_OBJ_NULL ) {
349400 nlr_raise (mp_obj_new_exception (& mp_type_KeyError ));
@@ -353,7 +404,7 @@ STATIC mp_obj_t set_remove(mp_obj_t self_in, mp_obj_t item) {
353404STATIC MP_DEFINE_CONST_FUN_OBJ_2 (set_remove_obj , set_remove );
354405
355406STATIC mp_obj_t set_symmetric_difference_update (mp_obj_t self_in , mp_obj_t other_in ) {
356- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
407+ check_set ( self_in );
357408 mp_obj_set_t * self = self_in ;
358409 mp_obj_t iter = mp_getiter (other_in );
359410 mp_obj_t next ;
@@ -365,10 +416,11 @@ STATIC mp_obj_t set_symmetric_difference_update(mp_obj_t self_in, mp_obj_t other
365416STATIC MP_DEFINE_CONST_FUN_OBJ_2 (set_symmetric_difference_update_obj , set_symmetric_difference_update );
366417
367418STATIC mp_obj_t set_symmetric_difference (mp_obj_t self_in , mp_obj_t other_in ) {
368- assert (MP_OBJ_IS_TYPE (self_in , & mp_type_set ));
369- self_in = set_copy (self_in );
370- set_symmetric_difference_update (self_in , other_in );
371- return self_in ;
419+ check_set_or_frozenset (self_in );
420+ mp_obj_set_t * self_out = set_copy_as_mutable (self_in );
421+ set_symmetric_difference_update (self_out , other_in );
422+ self_out -> base .type = ((mp_obj_set_t * )self_in )-> base .type ;
423+ return self_out ;
372424}
373425STATIC MP_DEFINE_CONST_FUN_OBJ_2 (set_symmetric_difference_obj , set_symmetric_difference );
374426
@@ -382,7 +434,6 @@ STATIC void set_update_int(mp_obj_set_t *self, mp_obj_t other_in) {
382434
383435STATIC mp_obj_t set_update (uint n_args , const mp_obj_t * args ) {
384436 assert (n_args > 0 );
385- assert (MP_OBJ_IS_TYPE (args [0 ], & mp_type_set ));
386437
387438 for (int i = 1 ; i < n_args ; i ++ ) {
388439 set_update_int (args [0 ], args [i ]);
@@ -393,7 +444,7 @@ STATIC mp_obj_t set_update(uint n_args, const mp_obj_t *args) {
393444STATIC MP_DEFINE_CONST_FUN_OBJ_VAR (set_update_obj , 1 , set_update );
394445
395446STATIC mp_obj_t set_union (mp_obj_t self_in , mp_obj_t other_in ) {
396- assert ( MP_OBJ_IS_TYPE ( self_in , & mp_type_set ) );
447+ check_set_or_frozenset ( self_in );
397448 mp_obj_set_t * self = set_copy (self_in );
398449 set_update_int (self , other_in );
399450 return self ;
@@ -486,6 +537,19 @@ const mp_obj_type_t mp_type_set = {
486537 .locals_dict = (mp_obj_t )& set_locals_dict ,
487538};
488539
540+ #if MICROPY_ENABLE_FROZENSET
541+ const mp_obj_type_t mp_type_frozenset = {
542+ { & mp_type_type },
543+ .name = MP_QSTR_frozenset ,
544+ .print = set_print ,
545+ .make_new = set_make_new ,
546+ .unary_op = set_unary_op ,
547+ .binary_op = set_binary_op ,
548+ .getiter = set_getiter ,
549+ .locals_dict = (mp_obj_t )& set_locals_dict ,
550+ };
551+ #endif
552+
489553mp_obj_t mp_obj_new_set (int n_args , mp_obj_t * items ) {
490554 mp_obj_set_t * o = m_new_obj (mp_obj_set_t );
491555 o -> base .type = & mp_type_set ;
0 commit comments