Merge pull request RustPython#1128 from RustPython/re-kung-fu

windelbouwman · web-flow · commit 98e6486c9863 · 2019-07-10T20:06:43.000+02:00
Add time.gmtime function.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/tests/snippets/test_re.py b/tests/snippets/test_re.py
@@ -26,8 +26,3 @@
 # TODO:
 # assert mo.group(0) == '_boe0'
 
-from string import Template
-s = Template('$who likes $what')
-# TODO:
-# r = s.substitute(who='tim', what='kung pow')
-# print(r)
diff --git a/tests/snippets/test_string.py b/tests/snippets/test_string.py
@@ -11,3 +11,15 @@
 # FIXME
 #assert string.whitespace == ' \t\n\r\x0b\x0c', string.whitespace
 #assert string.printable == '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ \t\n\r\x0b\x0c'
+
+assert string.capwords('bla bla', ' ') == 'Bla Bla'
+
+from string import Template
+s = Template('$who likes $what')
+# TODO:
+# r = s.substitute(who='tim', what='kung pow')
+# print(r)
+
+from string import Formatter
+
+f = Formatter()
diff --git a/tests/snippets/test_time.py b/tests/snippets/test_time.py
@@ -0,0 +1,21 @@
+
+
+import time
+
+x = time.gmtime(1000)
+
+assert x.tm_year == 1970
+assert x.tm_min == 16
+assert x.tm_sec == 40
+
+s = time.strftime('%Y-%m-%d-%H-%M-%S', x)
+# print(s)
+assert s == '1970-01-01-00-16-40'
+
+x2 = time.strptime(s, '%Y-%m-%d-%H-%M-%S')
+assert x2.tm_min == 16
+
+s = time.asctime(x)
+# print(s)
+assert s == 'Thu Jan  1 00:16:40 1970'
+
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -40,6 +40,7 @@ regex = "1"
 rustc_version_runtime = "0.1.*"
 statrs = "0.10.0"
 caseless = "0.2.1"
+chrono = "0.4.7"
 unicode-segmentation = "1.2.1"
 unicode-xid = "0.1.0"
 lazy_static = "^1.0.1"
diff --git a/vm/src/stdlib/time_module.rs b/vm/src/stdlib/time_module.rs
@@ -1,13 +1,22 @@
 //! The python `time` module.
-
+/// See also:
+/// https://docs.python.org/3/library/time.html
+use std::fmt;
 use std::thread;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
-use crate::function::PyFuncArgs;
-use crate::obj::objfloat;
-use crate::pyobject::{PyObjectRef, PyResult};
+use crate::function::{OptionalArg, PyFuncArgs};
+use crate::obj::objstr::PyStringRef;
+use crate::obj::objtype::PyClassRef;
+use crate::obj::{objfloat, objint, objtype};
+use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue};
 use crate::vm::VirtualMachine;
 
+use num_traits::cast::ToPrimitive;
+
+use chrono::naive::NaiveDateTime;
+use chrono::{Datelike, Timelike};
+
 fn time_sleep(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(seconds, Some(vm.ctx.float_type()))]);
     let seconds = objfloat::get_value(seconds);
@@ -32,11 +41,209 @@ fn time_time(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(value)
 }
 
+fn time_monotonic(vm: &VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args);
+    // TODO: implement proper monotonic time!
+    let x = match SystemTime::now().duration_since(UNIX_EPOCH) {
+        Ok(v) => duration_to_f64(v),
+        Err(err) => panic!("Error: {:?}", err),
+    };
+    let value = vm.ctx.new_float(x);
+    Ok(value)
+}
+
+fn pyfloat_to_secs_and_nanos(seconds: &PyObjectRef) -> (i64, u32) {
+    let seconds = objfloat::get_value(seconds);
+    let secs: i64 = seconds.trunc() as i64;
+    let nanos: u32 = (seconds.fract() * 1e9) as u32;
+    (secs, nanos)
+}
+
+fn pyobj_to_naive_date_time(
+    value: &PyObjectRef,
+    vm: &VirtualMachine,
+) -> PyResult<Option<NaiveDateTime>> {
+    if objtype::isinstance(value, &vm.ctx.float_type()) {
+        let (seconds, nanos) = pyfloat_to_secs_and_nanos(&value);
+        let dt = NaiveDateTime::from_timestamp(seconds, nanos);
+        Ok(Some(dt))
+    } else if objtype::isinstance(&value, &vm.ctx.int_type()) {
+        let seconds = objint::get_value(&value).to_i64().unwrap();
+        let dt = NaiveDateTime::from_timestamp(seconds, 0);
+        Ok(Some(dt))
+    } else {
+        Err(vm.new_type_error("Expected float, int or None".to_string()))
+    }
+}
+
+/// https://docs.python.org/3/library/time.html?highlight=gmtime#time.gmtime
+fn time_gmtime(secs: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<PyStructTime> {
+    let default = chrono::offset::Utc::now().naive_utc();
+    let instant = match secs {
+        OptionalArg::Present(secs) => pyobj_to_naive_date_time(&secs, vm)?.unwrap_or(default),
+        OptionalArg::Missing => default,
+    };
+    let value = PyStructTime::new(instant);
+    Ok(value)
+}
+
+fn time_localtime(secs: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<PyStructTime> {
+    let instant = optional_or_localtime(secs, vm)?;
+    let value = PyStructTime::new(instant);
+    Ok(value)
+}
+
+/// Construct a localtime from the optional seconds, or get the current local time.
+fn optional_or_localtime(
+    secs: OptionalArg<PyObjectRef>,
+    vm: &VirtualMachine,
+) -> PyResult<NaiveDateTime> {
+    let default = chrono::offset::Local::now().naive_local();
+    let instant = match secs {
+        OptionalArg::Present(secs) => pyobj_to_naive_date_time(&secs, vm)?.unwrap_or(default),
+        OptionalArg::Missing => default,
+    };
+    Ok(instant)
+}
+
+const CFMT: &str = "%a %b %e %H:%M:%S %Y";
+
+fn time_asctime(t: OptionalArg<PyStructTimeRef>, vm: &VirtualMachine) -> PyResult {
+    let default = chrono::offset::Local::now().naive_local();
+    let instant = match t {
+        OptionalArg::Present(t) => t.get_date_time(),
+        OptionalArg::Missing => default,
+    };
+    let formatted_time = instant.format(&CFMT).to_string();
+    Ok(vm.ctx.new_str(formatted_time))
+}
+
+fn time_ctime(secs: OptionalArg<PyObjectRef>, vm: &VirtualMachine) -> PyResult<String> {
+    let instant = optional_or_localtime(secs, vm)?;
+    let formatted_time = instant.format(&CFMT).to_string();
+    Ok(formatted_time)
+}
+
+fn time_strftime(
+    format: PyStringRef,
+    t: OptionalArg<PyStructTimeRef>,
+    vm: &VirtualMachine,
+) -> PyResult {
+    let default = chrono::offset::Local::now().naive_local();
+    let instant = match t {
+        OptionalArg::Present(t) => t.get_date_time(),
+        OptionalArg::Missing => default,
+    };
+    let formatted_time = instant.format(&format.value).to_string();
+    Ok(vm.ctx.new_str(formatted_time))
+}
+
+fn time_strptime(
+    string: PyStringRef,
+    format: OptionalArg<PyStringRef>,
+    vm: &VirtualMachine,
+) -> PyResult<PyStructTime> {
+    let format: String = match format {
+        OptionalArg::Present(format) => format.value.clone(),
+        OptionalArg::Missing => "%a %b %H:%M:%S %Y".to_string(),
+    };
+    let instant = NaiveDateTime::parse_from_str(&string.value, &format)
+        .map_err(|e| vm.new_value_error(format!("Parse error: {:?}", e)))?;
+    let struct_time = PyStructTime::new(instant);
+    Ok(struct_time)
+}
+
+#[pyclass(name = "struct_time")]
+struct PyStructTime {
+    tm: NaiveDateTime,
+}
+
+type PyStructTimeRef = PyRef<PyStructTime>;
+
+impl fmt::Debug for PyStructTime {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "struct_time()")
+    }
+}
+
+impl PyValue for PyStructTime {
+    fn class(vm: &VirtualMachine) -> PyClassRef {
+        vm.class("time", "struct_time")
+    }
+}
+
+#[pyimpl]
+impl PyStructTime {
+    fn new(tm: NaiveDateTime) -> Self {
+        PyStructTime { tm }
+    }
+
+    #[pymethod(name = "__repr__")]
+    fn repr(&self, _vm: &VirtualMachine) -> String {
+        // TODO: extract year day and isdst somehow..
+        format!(
+            "time.struct_time(tm_year={}, tm_mon={}, tm_mday={}, tm_hour={}, tm_min={}, tm_sec={}, tm_wday={})",
+            self.tm.date().year(), self.tm.date().month(), self.tm.date().day(),
+            self.tm.time().hour(), self.tm.time().minute(), self.tm.time().second(),
+            self.tm.date().weekday().num_days_from_monday()
+            )
+    }
+
+    fn get_date_time(&self) -> NaiveDateTime {
+        self.tm
+    }
+
+    #[pyproperty(name = "tm_year")]
+    fn tm_year(&self, _vm: &VirtualMachine) -> i32 {
+        self.tm.date().year()
+    }
+
+    #[pyproperty(name = "tm_mon")]
+    fn tm_mon(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.date().month()
+    }
+
+    #[pyproperty(name = "tm_mday")]
+    fn tm_mday(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.date().day()
+    }
+
+    #[pyproperty(name = "tm_hour")]
+    fn tm_hour(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.time().hour()
+    }
+
+    #[pyproperty(name = "tm_min")]
+    fn tm_min(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.time().minute()
+    }
+
+    #[pyproperty(name = "tm_sec")]
+    fn tm_sec(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.time().second()
+    }
+
+    #[pyproperty(name = "tm_wday")]
+    fn tm_wday(&self, _vm: &VirtualMachine) -> u32 {
+        self.tm.date().weekday().num_days_from_monday()
+    }
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
+    let struct_time_type = PyStructTime::make_class(ctx);
+
     py_module!(vm, "time", {
+        "asctime" => ctx.new_rustfunc(time_asctime),
+        "ctime" => ctx.new_rustfunc(time_ctime),
+        "gmtime" => ctx.new_rustfunc(time_gmtime),
+        "localtime" => ctx.new_rustfunc(time_localtime),
+        "monotonic" => ctx.new_rustfunc(time_monotonic),
+        "strftime" => ctx.new_rustfunc(time_strftime),
+        "strptime" => ctx.new_rustfunc(time_strptime),
         "sleep" => ctx.new_rustfunc(time_sleep),
+        "struct_time" => struct_time_type,
         "time" => ctx.new_rustfunc(time_time)
     })
 }
