Implement Add<Months> and Sub<Months> for NaiveDate (#731)

* Add add_months()
* Months struct with Add and Sub impls
* Update changelog

.gitignore

@@ -1,3 +1,6 @@
 target
 Cargo.lock
 .tool-versions
+
+# for jetbrains users
+.idea/

CHANGELOG.md

@@ -36,6 +36,7 @@ Versions with only mechanical changes will be omitted from the following list.
 * Fix the behavior of `Duration::abs()` for negative durations with non-zero nanos
 * Add compatibility with rfc2822 comments (#733)
 * Make `js-sys` and `wasm-bindgen` enabled by default when target is `wasm32-unknown-unknown` for ease of API discovery
+* Add the `Months` struct and associated `Add` and `Sub` impls
 
 ## 0.4.19
 

src/lib.rs

@@ -507,7 +507,7 @@ mod weekday;
 pub use weekday::{ParseWeekdayError, Weekday};
 
 mod month;
-pub use month::{Month, ParseMonthError};
+pub use month::{Month, Months, ParseMonthError};
 
 mod traits;
 pub use traits::{Datelike, Timelike};

src/month.rs

@@ -189,6 +189,10 @@ impl num_traits::FromPrimitive for Month {
     }
 }
 
+/// A duration in calendar months
+#[derive(Clone, Debug, PartialEq)]
+pub struct Months(pub usize);
+
 /// An error resulting from reading `<Month>` value with `FromStr`.
 #[derive(Clone, PartialEq)]
 pub struct ParseMonthError {

src/naive/date.rs

@@ -17,6 +17,7 @@ use rkyv::{Archive, Deserialize, Serialize};
 use crate::format::DelayedFormat;
 use crate::format::{parse, ParseError, ParseResult, Parsed, StrftimeItems};
 use crate::format::{Item, Numeric, Pad};
+use crate::month::Months;
 use crate::naive::{IsoWeek, NaiveDateTime, NaiveTime};
 use crate::oldtime::Duration as OldDuration;
 use crate::{Datelike, Duration, Weekday};
@@ -596,6 +597,33 @@ impl NaiveDate {
         parsed.to_naive_date()
     }
 
+    /// Private function to calculate necessary primitives for `Add<Month>`
+    ///
+    /// # Arguments
+    ///
+    /// * `delta` - Number of months (+/-) to add
+    ///
+    /// # Returns
+    ///
+    /// A new NaiveDate on the first day of the resulting year & month
+    fn add_months_get_first_day(&self, delta: i32) -> NaiveDate {
+        let zeroed_months = self.month() as i32 - 1; // zero-based for modulo operations
+        let res_months = zeroed_months + delta;
+        let delta_years = if res_months < 0 {
+            if (-res_months) % 12 > 0 {
+                res_months / 12 - 1
+            } else {
+                res_months / 12
+            }
+        } else {
+            res_months / 12
+        };
+        let res_years = self.year() + delta_years;
+        let res_months = res_months % 12;
+        let res_months = if res_months < 0 { res_months + 12 } else { res_months };
+        NaiveDate::from_ymd(res_years, res_months as u32 + 1, 1)
+    }
+
     /// Makes a new `NaiveDateTime` from the current date and given `NaiveTime`.
     ///
     /// # Example
@@ -1561,6 +1589,59 @@ impl AddAssign<OldDuration> for NaiveDate {
     }
 }
 
+impl Add<Months> for NaiveDate {
+    type Output = NaiveDate;
+
+    /// An addition of months to `NaiveDate` clamped to valid days in resulting month.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use chrono::{Duration, NaiveDate, Months};
+    ///
+    /// let from_ymd = NaiveDate::from_ymd;
+    ///
+    /// assert_eq!(from_ymd(2014, 1, 1) + Months(1), from_ymd(2014, 2, 1));
+    /// assert_eq!(from_ymd(2014, 1, 1) + Months(11), from_ymd(2014, 12, 1));
+    /// assert_eq!(from_ymd(2014, 1, 1) + Months(12), from_ymd(2015, 1, 1));
+    /// assert_eq!(from_ymd(2014, 1, 1) + Months(13), from_ymd(2015, 2, 1));
+    /// assert_eq!(from_ymd(2014, 1, 31) + Months(1), from_ymd(2014, 2, 28));
+    /// assert_eq!(from_ymd(2020, 1, 31) + Months(1), from_ymd(2020, 2, 29));
+    /// ```
+    fn add(self, months: Months) -> Self::Output {
+        let target = self.add_months_get_first_day(months.0 as i32);
+        let target_plus = target.add_months_get_first_day(1);
+        let last_day = target_plus.sub(Duration::days(1));
+        let day = core::cmp::min(self.day(), last_day.day());
+        NaiveDate::from_ymd(target.year(), target.month(), day)
+    }
+}
+
+impl Sub<Months> for NaiveDate {
+    type Output = NaiveDate;
+
+    /// A subtraction of Months from `NaiveDate` clamped to valid days in resulting month.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use chrono::{Duration, NaiveDate, Months};
+    ///
+    /// let from_ymd = NaiveDate::from_ymd;
+    ///
+    /// assert_eq!(from_ymd(2014, 1, 1) - Months(11), from_ymd(2013, 2, 1));
+    /// assert_eq!(from_ymd(2014, 1, 1) - Months(12), from_ymd(2013, 1, 1));
+    /// assert_eq!(from_ymd(2014, 1, 1) - Months(13), from_ymd(2012, 12, 1));
+    /// ```
+    fn sub(self, months: Months) -> Self::Output {
+        let target = self.add_months_get_first_day(-(months.0 as i32));
+        let target_plus = target.add_months_get_first_day(1);
+        let last_day = target_plus.sub(Duration::days(1));
+        let day = core::cmp::min(self.day(), last_day.day());
+        NaiveDate::from_ymd(target.year(), target.month(), day)
+    }
+}
+
 /// A subtraction of `Duration` from `NaiveDate` discards the fractional days,
 /// rounding to the closest integral number of days towards `Duration::zero()`.
 /// It is the same as the addition with a negated `Duration`.
@@ -2002,6 +2083,46 @@ mod tests {
     use crate::{Datelike, Weekday};
     use std::{i32, u32};
 
+    #[test]
+    fn test_add_months_get_first_day() {
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 1).add_months_get_first_day(1),
+            NaiveDate::from_ymd(2014, 2, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 31).add_months_get_first_day(1),
+            NaiveDate::from_ymd(2014, 2, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2020, 1, 10).add_months_get_first_day(1),
+            NaiveDate::from_ymd(2020, 2, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 1).add_months_get_first_day(-1),
+            NaiveDate::from_ymd(2013, 12, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 31).add_months_get_first_day(-1),
+            NaiveDate::from_ymd(2013, 12, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2020, 1, 10).add_months_get_first_day(-1),
+            NaiveDate::from_ymd(2019, 12, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 10).add_months_get_first_day(-11),
+            NaiveDate::from_ymd(2013, 2, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 10).add_months_get_first_day(-12),
+            NaiveDate::from_ymd(2013, 1, 1)
+        );
+        assert_eq!(
+            NaiveDate::from_ymd(2014, 1, 10).add_months_get_first_day(-13),
+            NaiveDate::from_ymd(2012, 12, 1)
+        );
+    }
+
     #[test]
     fn test_readme_doomsday() {
         use num_iter::range_inclusive;