@@ -21,6 +21,9 @@ use itertools::Itertools;
2121use malachite_bigint:: BigInt ;
2222use num_traits:: ToPrimitive ;
2323
24+ const STRING_WITHOUT_ENCODING : & str = "string argument without an encoding" ;
25+ const ENCODING_WITHOUT_STRING : & str = "encoding without a string argument" ;
26+
2427#[ derive( Debug , Default , Clone ) ]
2528pub struct PyBytesInner {
2629 pub ( super ) elements : Vec < u8 > ,
@@ -75,6 +78,18 @@ impl ByteInnerNewOptions {
7578 Ok ( vec ! [ 0 ; size] . into ( ) )
7679 }
7780
81+ fn handle_object_fallback ( obj : PyObjectRef , vm : & VirtualMachine ) -> PyResult < PyBytesInner > {
82+ match_class ! ( match obj {
83+ i @ PyInt => {
84+ Self :: get_value_from_size( i, vm)
85+ }
86+ _s @ PyStr => Err ( vm. new_type_error( STRING_WITHOUT_ENCODING . to_owned( ) ) ) ,
87+ obj => {
88+ Self :: get_value_from_source( obj, vm)
89+ }
90+ } )
91+ }
92+
7893 pub fn get_bytes ( self , cls : PyTypeRef , vm : & VirtualMachine ) -> PyResult < PyBytesRef > {
7994 let inner = match ( & self . source , & self . encoding , & self . errors ) {
8095 ( OptionalArg :: Present ( obj) , OptionalArg :: Missing , OptionalArg :: Missing ) => {
@@ -113,40 +128,48 @@ impl ByteInnerNewOptions {
113128 }
114129
115130 pub fn get_bytearray_inner ( self , vm : & VirtualMachine ) -> PyResult < PyBytesInner > {
116- const STRING_WITHOUT_ENCODING : & str = "string argument without an encoding" ;
117- const ENCODING_WITHOUT_STRING : & str = "encoding without a string argument" ;
118-
119131 match ( self . source , self . encoding , self . errors ) {
120132 ( OptionalArg :: Present ( obj) , OptionalArg :: Missing , OptionalArg :: Missing ) => {
121- match_class ! ( match obj {
122- i @ PyInt => {
123- Ok ( Self :: get_value_from_size( i, vm) ?)
124- }
125- _s @ PyStr => Err ( STRING_WITHOUT_ENCODING ) ,
126- obj => {
127- Ok ( Self :: get_value_from_source( obj, vm) ?)
133+ // Try __index__ first to handle int-like objects that might raise custom exceptions
134+ if let Some ( index_result) = obj. try_index_opt ( vm) {
135+ match index_result {
136+ Ok ( index) => Self :: get_value_from_size ( index, vm) ,
137+ Err ( e) => {
138+ // Only propagate non-TypeError exceptions
139+ // TypeError means the object doesn't support __index__, so fall back
140+ if e. fast_isinstance ( vm. ctx . exceptions . type_error ) {
141+ // Fall back to treating as buffer-like object
142+ Self :: handle_object_fallback ( obj, vm)
143+ } else {
144+ // Propagate other exceptions (e.g., ZeroDivisionError)
145+ Err ( e)
146+ }
147+ }
128148 }
129- } )
149+ } else {
150+ Self :: handle_object_fallback ( obj, vm)
151+ }
130152 }
131153 ( OptionalArg :: Present ( obj) , OptionalArg :: Present ( encoding) , errors) => {
132154 if let Ok ( s) = obj. downcast :: < PyStr > ( ) {
133- Ok ( Self :: get_value_from_string ( s, encoding, errors, vm) ? )
155+ Self :: get_value_from_string ( s, encoding, errors, vm)
134156 } else {
135- Err ( ENCODING_WITHOUT_STRING )
157+ Err ( vm . new_type_error ( ENCODING_WITHOUT_STRING . to_owned ( ) ) )
136158 }
137159 }
138160 ( OptionalArg :: Missing , OptionalArg :: Missing , OptionalArg :: Missing ) => {
139161 Ok ( PyBytesInner :: default ( ) )
140162 }
141- ( OptionalArg :: Missing , OptionalArg :: Present ( _) , _) => Err ( ENCODING_WITHOUT_STRING ) ,
163+ ( OptionalArg :: Missing , OptionalArg :: Present ( _) , _) => {
164+ Err ( vm. new_type_error ( ENCODING_WITHOUT_STRING . to_owned ( ) ) )
165+ }
142166 ( OptionalArg :: Missing , _, OptionalArg :: Present ( _) ) => {
143- Err ( "errors without a string argument" )
167+ Err ( vm . new_type_error ( "errors without a string argument" . to_owned ( ) ) )
144168 }
145169 ( OptionalArg :: Present ( _) , OptionalArg :: Missing , OptionalArg :: Present ( _) ) => {
146- Err ( STRING_WITHOUT_ENCODING )
170+ Err ( vm . new_type_error ( STRING_WITHOUT_ENCODING . to_owned ( ) ) )
147171 }
148172 }
149- . map_err ( |e| vm. new_type_error ( e. to_owned ( ) ) )
150173 }
151174}
152175
0 commit comments