@@ -338,7 +338,10 @@ mod decl {
338338 } ;
339339 use libc:: pollfd;
340340 use num_traits:: { Signed , ToPrimitive } ;
341- use std:: time:: { Duration , Instant } ;
341+ use std:: {
342+ convert:: TryFrom ,
343+ time:: { Duration , Instant } ,
344+ } ;
342345
343346 #[ derive( Default ) ]
344347 pub ( super ) struct TimeoutArg < const MILLIS : bool > ( pub Option < Duration > ) ;
@@ -417,25 +420,62 @@ mod decl {
417420 search ( fds, fd) . ok ( ) . map ( |i| fds. remove ( i) )
418421 }
419422
423+ // new EventMask type
424+ #[ derive( Copy , Clone ) ]
425+ #[ repr( transparent) ]
426+ pub struct EventMask ( pub i16 ) ;
427+
428+ impl TryFromObject for EventMask {
429+ fn try_from_object ( vm : & VirtualMachine , obj : PyObjectRef ) -> PyResult < Self > {
430+ use crate :: builtins:: PyInt ;
431+ let int = obj
432+ . downcast :: < PyInt > ( )
433+ . map_err ( |_| vm. new_type_error ( "argument must be an integer" . to_owned ( ) ) ) ?;
434+
435+ let val = int. as_bigint ( ) ;
436+ if val. is_negative ( ) {
437+ return Err ( vm. new_value_error ( "negative event mask" . to_owned ( ) ) ) ;
438+ }
439+
440+ // Try converting to i16, should raise OverflowError if too large
441+ let mask = i16:: try_from ( val) . map_err ( |_| {
442+ vm. new_overflow_error ( "event mask value out of range" . to_owned ( ) )
443+ } ) ?;
444+
445+ Ok ( EventMask ( mask) )
446+ }
447+ }
448+
420449 const DEFAULT_EVENTS : i16 = libc:: POLLIN | libc:: POLLPRI | libc:: POLLOUT ;
421450
422451 #[ pyclass]
423452 impl PyPoll {
424453 #[ pymethod]
425- fn register ( & self , Fildes ( fd) : Fildes , eventmask : OptionalArg < u16 > ) {
426- insert_fd (
427- & mut self . fds . lock ( ) ,
428- fd,
429- eventmask. map_or ( DEFAULT_EVENTS , |e| e as i16 ) ,
430- )
454+ fn register (
455+ & self ,
456+ Fildes ( fd) : Fildes ,
457+ eventmask : OptionalArg < EventMask > ,
458+ ) -> PyResult < ( ) > {
459+ let mask = match eventmask {
460+ OptionalArg :: Present ( event_mask) => event_mask. 0 ,
461+ OptionalArg :: Missing => DEFAULT_EVENTS ,
462+ } ;
463+ insert_fd ( & mut self . fds . lock ( ) , fd, mask) ;
464+ Ok ( ( ) )
431465 }
432466
433467 #[ pymethod]
434- fn modify ( & self , Fildes ( fd) : Fildes , eventmask : u16 ) -> io:: Result < ( ) > {
468+ fn modify (
469+ & self ,
470+ Fildes ( fd) : Fildes ,
471+ eventmask : EventMask ,
472+ vm : & VirtualMachine ,
473+ ) -> PyResult < ( ) > {
435474 let mut fds = self . fds . lock ( ) ;
475+ // CPython raises KeyError if fd is not registered, match that behavior
436476 let pfd = get_fd_mut ( & mut fds, fd)
437- . ok_or_else ( || io :: Error :: from_raw_os_error ( libc :: ENOENT ) ) ?;
438- pfd. events = eventmask as i16 ;
477+ . ok_or_else ( || vm . new_key_error ( vm . ctx . new_int ( fd ) . into ( ) ) ) ?;
478+ pfd. events = eventmask. 0 ;
439479 Ok ( ( ) )
440480 }
441481
0 commit comments