Implement operations for core::time::Duration on DateTime types

chronotope · Aug 28, 2023 · 169409e · 169409e
1 parent 1a3c43a
commit 169409e
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 2 deletions.
diff --git a/src/datetime/mod.rs b/src/datetime/mod.rs
@@ -12,6 +12,7 @@ use core::borrow::Borrow;
 use core::cmp::Ordering;
 use core::fmt::Write;
 use core::ops::{Add, AddAssign, Sub, SubAssign};
+use core::time::Duration;
 use core::{fmt, hash, str};
 #[cfg(feature = "std")]
 use std::string::ToString;
@@ -1160,6 +1161,17 @@ impl<Tz: TimeZone> Add<OldDuration> for DateTime<Tz> {
     }
 }
 
+impl<Tz: TimeZone> Add<Duration> for DateTime<Tz> {
+    type Output = DateTime<Tz>;
+
+    #[inline]
+    fn add(self, rhs: Duration) -> DateTime<Tz> {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        self.checked_add_signed(rhs).expect("`DateTime + Duration` overflowed")
+    }
+}
+
 impl<Tz: TimeZone> AddAssign<OldDuration> for DateTime<Tz> {
     #[inline]
     fn add_assign(&mut self, rhs: OldDuration) {
@@ -1170,6 +1182,15 @@ impl<Tz: TimeZone> AddAssign<OldDuration> for DateTime<Tz> {
     }
 }
 
+impl<Tz: TimeZone> AddAssign<Duration> for DateTime<Tz> {
+    #[inline]
+    fn add_assign(&mut self, rhs: Duration) {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        *self += rhs;
+    }
+}
+
 impl<Tz: TimeZone> Add<Months> for DateTime<Tz> {
     type Output = DateTime<Tz>;
 
@@ -1187,6 +1208,17 @@ impl<Tz: TimeZone> Sub<OldDuration> for DateTime<Tz> {
     }
 }
 
+impl<Tz: TimeZone> Sub<Duration> for DateTime<Tz> {
+    type Output = DateTime<Tz>;
+
+    #[inline]
+    fn sub(self, rhs: Duration) -> DateTime<Tz> {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        self.checked_sub_signed(rhs).expect("`DateTime - Duration` overflowed")
+    }
+}
+
 impl<Tz: TimeZone> SubAssign<OldDuration> for DateTime<Tz> {
     #[inline]
     fn sub_assign(&mut self, rhs: OldDuration) {
@@ -1197,6 +1229,18 @@ impl<Tz: TimeZone> SubAssign<OldDuration> for DateTime<Tz> {
     }
 }
 
+impl<Tz: TimeZone> SubAssign<Duration> for DateTime<Tz> {
+    #[inline]
+    fn sub_assign(&mut self, rhs: Duration) {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        let datetime =
+            self.datetime.checked_sub_signed(rhs).expect("`DateTime - Duration` overflowed");
+        let tz = self.timezone();
+        *self = tz.from_utc_datetime(&datetime)
+    }
+}
+
 impl<Tz: TimeZone> Sub<Months> for DateTime<Tz> {
     type Output = DateTime<Tz>;
 
@@ -1334,8 +1378,6 @@ impl From<SystemTime> for DateTime<Local> {
 #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
 impl<Tz: TimeZone> From<DateTime<Tz>> for SystemTime {
     fn from(dt: DateTime<Tz>) -> SystemTime {
-        use std::time::Duration;
-
         let sec = dt.timestamp();
         let nsec = dt.timestamp_subsec_nanos();
         if sec < 0 {

diff --git a/src/naive/datetime/mod.rs b/src/naive/datetime/mod.rs
@@ -7,6 +7,7 @@
 use core::borrow::Borrow;
 use core::fmt::Write;
 use core::ops::{Add, AddAssign, Sub, SubAssign};
+use core::time::Duration;
 use core::{fmt, str};
 
 #[cfg(feature = "rkyv")]
@@ -1516,13 +1517,31 @@ impl Add<OldDuration> for NaiveDateTime {
     }
 }
 
+impl Add<Duration> for NaiveDateTime {
+    type Output = NaiveDateTime;
+
+    #[inline]
+    fn add(self, rhs: Duration) -> NaiveDateTime {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        self.checked_add_signed(rhs).expect("`NaiveDateTime + Duration` overflowed")
+    }
+}
+
 impl AddAssign<OldDuration> for NaiveDateTime {
     #[inline]
     fn add_assign(&mut self, rhs: OldDuration) {
         *self = self.add(rhs);
     }
 }
 
+impl AddAssign<Duration> for NaiveDateTime {
+    #[inline]
+    fn add_assign(&mut self, rhs: Duration) {
+        *self = self.add(rhs);
+    }
+}
+
 impl Add<Months> for NaiveDateTime {
     type Output = NaiveDateTime;
 
@@ -1625,13 +1644,31 @@ impl Sub<OldDuration> for NaiveDateTime {
     }
 }
 
+impl Sub<Duration> for NaiveDateTime {
+    type Output = NaiveDateTime;
+
+    #[inline]
+    fn sub(self, rhs: Duration) -> NaiveDateTime {
+        let rhs = OldDuration::from_std(rhs)
+            .expect("overflow converting from core::time::Duration to chrono::Duration");
+        self.checked_sub_signed(rhs).expect("`NaiveDateTime - Duration` overflowed")
+    }
+}
+
 impl SubAssign<OldDuration> for NaiveDateTime {
     #[inline]
     fn sub_assign(&mut self, rhs: OldDuration) {
         *self = self.sub(rhs);
     }
 }
 
+impl SubAssign<Duration> for NaiveDateTime {
+    #[inline]
+    fn sub_assign(&mut self, rhs: Duration) {
+        *self = self.sub(rhs);
+    }
+}
+
 /// A subtraction of Months from `NaiveDateTime` clamped to valid days in resulting month.
 ///
 /// # Panics