@@ -2,7 +2,7 @@ use super::{PyMethod, VirtualMachine};
22use crate :: {
33 builtins:: { PyInt , PyIntRef , PyStr , PyStrRef } ,
44 object:: { AsObject , PyObject , PyObjectRef , PyResult } ,
5- protocol:: { PyIterReturn , PyNumberBinaryOp , PySequence } ,
5+ protocol:: { PyIterReturn , PyNumberBinaryOp , PyNumberTernaryOp , PySequence } ,
66 types:: PyComparisonOp ,
77} ;
88use num_traits:: ToPrimitive ;
@@ -15,6 +15,14 @@ macro_rules! binary_func {
1515 } ;
1616}
1717
18+ macro_rules! ternary_func {
19+ ( $fn: ident, $op_slot: ident, $op: expr) => {
20+ pub fn $fn( & self , a: & PyObject , b: & PyObject , c: & PyObject ) -> PyResult {
21+ self . ternary_op( a, b, c, PyNumberTernaryOp :: $op_slot, $op)
22+ }
23+ } ;
24+ }
25+
1826macro_rules! inplace_binary_func {
1927 ( $fn: ident, $iop_slot: ident, $op_slot: ident, $op: expr) => {
2028 pub fn $fn( & self , a: & PyObject , b: & PyObject ) -> PyResult {
@@ -29,6 +37,21 @@ macro_rules! inplace_binary_func {
2937 } ;
3038}
3139
40+ macro_rules! inplace_ternary_func {
41+ ( $fn: ident, $iop_slot: ident, $op_slot: ident, $op: expr) => {
42+ pub fn $fn( & self , a: & PyObject , b: & PyObject , c: & PyObject ) -> PyResult {
43+ self . ternary_iop(
44+ a,
45+ b,
46+ c,
47+ PyNumberTernaryOp :: $iop_slot,
48+ PyNumberTernaryOp :: $op_slot,
49+ $op,
50+ )
51+ }
52+ } ;
53+ }
54+
3255/// Collection of operators
3356impl VirtualMachine {
3457 #[ inline]
@@ -229,10 +252,103 @@ impl VirtualMachine {
229252 Err ( self . new_unsupported_binop_error ( a, b, op) )
230253 }
231254
255+ fn ternary_op (
256+ & self ,
257+ a : & PyObject ,
258+ b : & PyObject ,
259+ c : & PyObject ,
260+ op_slot : PyNumberTernaryOp ,
261+ op_str : & str ,
262+ ) -> PyResult {
263+ let class_a = a. class ( ) ;
264+ let class_b = b. class ( ) ;
265+ let class_c = c. class ( ) ;
266+
267+ let slot_a = class_a. slots . as_number . left_ternary_op ( op_slot) ;
268+ let mut slot_b = None ;
269+
270+ if !class_a. is ( class_b) {
271+ let slot_bb = class_b. slots . as_number . right_ternary_op ( op_slot) ;
272+ if slot_bb. map ( |x| x as usize ) != slot_a. map ( |x| x as usize ) {
273+ slot_b = slot_bb;
274+ }
275+ }
276+
277+ if let Some ( slot_a) = slot_a {
278+ if let Some ( slot_bb) = slot_b {
279+ if class_b. fast_issubclass ( class_a) {
280+ let ret = slot_bb ( a, b, c, self ) ?;
281+ if !ret. is ( & self . ctx . not_implemented ) {
282+ return Ok ( ret) ;
283+ }
284+ slot_b = None ;
285+ }
286+ }
287+ let ret = slot_a ( a, b, c, self ) ?;
288+ if !ret. is ( & self . ctx . not_implemented ) {
289+ return Ok ( ret) ;
290+ }
291+ }
292+
293+ if let Some ( slot_b) = slot_b {
294+ let ret = slot_b ( a, b, c, self ) ?;
295+ if !ret. is ( & self . ctx . not_implemented ) {
296+ return Ok ( ret) ;
297+ }
298+ }
299+
300+ if let Some ( slot_c) = class_c. slots . as_number . left_ternary_op ( op_slot) {
301+ if slot_a. map_or ( false , |slot_a| ( slot_a as usize ) != ( slot_c as usize ) )
302+ && slot_b. map_or ( false , |slot_b| ( slot_b as usize ) != ( slot_c as usize ) )
303+ {
304+ let ret = slot_c ( a, b, c, self ) ?;
305+ if !ret. is ( & self . ctx . not_implemented ) {
306+ return Ok ( ret) ;
307+ }
308+ }
309+ }
310+
311+ Err ( if self . is_none ( c) {
312+ self . new_type_error ( format ! (
313+ "unsupported operand type(s) for {}: \
314+ ,
315+ op_str,
316+ a. class( ) ,
317+ b. class( )
318+ ) )
319+ } else {
320+ self . new_type_error ( format ! (
321+ "unsupported operand type(s) for {}: \
322+ ,
323+ op_str,
324+ a. class( ) ,
325+ b. class( ) ,
326+ c. class( )
327+ ) )
328+ } )
329+ }
330+
331+ fn ternary_iop (
332+ & self ,
333+ a : & PyObject ,
334+ b : & PyObject ,
335+ c : & PyObject ,
336+ iop_slot : PyNumberTernaryOp ,
337+ op_slot : PyNumberTernaryOp ,
338+ op_str : & str ,
339+ ) -> PyResult {
340+ if let Some ( slot) = a. class ( ) . slots . as_number . left_ternary_op ( iop_slot) {
341+ let x = slot ( a, b, c, self ) ?;
342+ if !x. is ( & self . ctx . not_implemented ) {
343+ return Ok ( x) ;
344+ }
345+ }
346+ self . ternary_op ( a, b, c, op_slot, op_str)
347+ }
348+
232349 binary_func ! ( _sub, Subtract , "-" ) ;
233350 binary_func ! ( _mod, Remainder , "%" ) ;
234351 binary_func ! ( _divmod, Divmod , "divmod" ) ;
235- binary_func ! ( _pow, Power , "**" ) ;
236352 binary_func ! ( _lshift, Lshift , "<<" ) ;
237353 binary_func ! ( _rshift, Rshift , ">>" ) ;
238354 binary_func ! ( _and, And , "&" ) ;
@@ -244,7 +360,6 @@ impl VirtualMachine {
244360
245361 inplace_binary_func ! ( _isub, InplaceSubtract , Subtract , "-=" ) ;
246362 inplace_binary_func ! ( _imod, InplaceRemainder , Remainder , "%=" ) ;
247- inplace_binary_func ! ( _ipow, InplacePower , Power , "**=" ) ;
248363 inplace_binary_func ! ( _ilshift, InplaceLshift , Lshift , "<<=" ) ;
249364 inplace_binary_func ! ( _irshift, InplaceRshift , Rshift , ">>=" ) ;
250365 inplace_binary_func ! ( _iand, InplaceAnd , And , "&=" ) ;
@@ -254,6 +369,9 @@ impl VirtualMachine {
254369 inplace_binary_func ! ( _itruediv, InplaceTrueDivide , TrueDivide , "/=" ) ;
255370 inplace_binary_func ! ( _imatmul, InplaceMatrixMultiply , MatrixMultiply , "@=" ) ;
256371
372+ ternary_func ! ( _pow, Power , "** or pow()" ) ;
373+ inplace_ternary_func ! ( _ipow, InplacePower , Power , "**=" ) ;
374+
257375 pub fn _add ( & self , a : & PyObject , b : & PyObject ) -> PyResult {
258376 let result = self . binary_op1 ( a, b, PyNumberBinaryOp :: Add ) ?;
259377 if !result. is ( & self . ctx . not_implemented ) {
0 commit comments