Add a TryFromObject impl for std::time::Duration

vm/src/pyobject.rs

@@ -1183,6 +1183,22 @@ pub trait PyClassImpl: PyClassDef {
     }
 }
 
+// TODO: find a better place to put this impl
+impl TryFromObject for std::time::Duration {
+    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
+        use std::time::Duration;
+        u64::try_from_object(vm, obj.clone())
+            .map(Duration::from_secs)
+            .or_else(|_| f64::try_from_object(vm, obj.clone()).map(Duration::from_secs_f64))
+            .map_err(|_| {
+                vm.new_type_error(format!(
+                    "expected an int or float for duration, got {}",
+                    obj.class()
+                ))
+            })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

vm/src/stdlib/socket.rs

@@ -280,20 +280,15 @@ impl PySocket {
     }
 
     #[pymethod]
-    fn settimeout(&self, timeout: Option<Either<f64, u64>>, vm: &VirtualMachine) -> PyResult<()> {
-        let mut block = if timeout.is_none() { Some(true) } else { None };
-        let timeout = timeout.and_then(|n| {
-            let dur = match n {
-                Either::A(f) => Duration::from_secs_f64(f),
-                Either::B(i) => Duration::from_secs(i),
-            };
-            if dur == Duration::from_secs(0) {
-                block = Some(false);
-                None
-            } else {
-                Some(dur)
-            }
-        });
+    fn settimeout(&self, timeout: Option<Duration>, vm: &VirtualMachine) -> PyResult<()> {
+        // timeout is None: blocking, no timeout
+        // timeout is 0: non-blocking, no timeout
+        // otherwise: timeout is timeout, don't change blocking
+        let (block, timeout) = match timeout {
+            None => (Some(true), None),
+            Some(d) if d == Duration::from_secs(0) => (Some(false), None),
+            Some(d) => (None, Some(d)),
+        };
         self.sock()
             .set_read_timeout(timeout)
             .map_err(|err| convert_sock_error(vm, err))?;

vm/src/stdlib/time_module.rs

@@ -15,9 +15,8 @@ use crate::pyobject::{Either, PyClassImpl, PyObjectRef, PyResult, TryFromObject}
 use crate::vm::VirtualMachine;
 
 #[cfg(unix)]
-fn time_sleep(seconds: f64, vm: &VirtualMachine) -> PyResult<()> {
-    // this is basically std::thread::sleep, but that catches interrupts and we don't want to
-    let dur = Duration::from_secs_f64(seconds);
+fn time_sleep(dur: Duration, vm: &VirtualMachine) -> PyResult<()> {
+    // this is basically std::thread::sleep, but that catches interrupts and we don't want to;
 
     let mut ts = libc::timespec {
         tv_sec: std::cmp::min(libc::time_t::max_value() as u64, dur.as_secs()) as libc::time_t,
@@ -34,8 +33,8 @@ fn time_sleep(seconds: f64, vm: &VirtualMachine) -> PyResult<()> {
 }
 
 #[cfg(not(unix))]
-fn time_sleep(seconds: f64, _vm: &VirtualMachine) {
-    std::thread::sleep(Duration::from_secs_f64(seconds));
+fn time_sleep(dur: Duration, _vm: &VirtualMachine) {
+    std::thread::sleep(dur);
 }
 
 #[cfg(any(not(target_arch = "wasm32"), target_os = "wasi"))]