/Users/andrewlamb/Software/arrow-rs/arrow-array/src/delta.rs

Source
// MIT License
//
// Copyright (c) 2020-2022 Oliver Margetts
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Copied from chronoutil crate

//! Contains utility functions for shifting Date objects.
use chrono::{DateTime, Datelike, Days, Months, TimeZone};
use std::cmp::Ordering;

/// Shift a date by the given number of months.
pub(crate) fn shift_months<D>(date: D, months: i32) -> D
where
    D: Datelike + std::ops::Add<Months, Output = D> + std::ops::Sub<Months, Output = D>,
{
    match months.cmp(&0) {
        Ordering::Equal => date,
        Ordering::Greater => date + Months::new(months as u32),
        Ordering::Less => date - Months::new(months.unsigned_abs()),
    }
}

/// Add the given number of months to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn add_months_datetime<Tz: TimeZone>(
    dt: DateTime<Tz>,
    months: i32,
) -> Option<DateTime<Tz>> {
    match months.cmp(&0) {
        Ordering::Equal => Some(dt),
        Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
        Ordering::Less => dt.checked_sub_months(Months::new(months.unsigned_abs())),
    }
}

/// Add the given number of days to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
    match days.cmp(&0) {
        Ordering::Equal => Some(dt),
        Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
        Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as u64)),
    }
}

/// Substract the given number of months to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn sub_months_datetime<Tz: TimeZone>(
    dt: DateTime<Tz>,
    months: i32,
) -> Option<DateTime<Tz>> {
    match months.cmp(&0) {
        Ordering::Equal => Some(dt),
        Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
        Ordering::Less => dt.checked_add_months(Months::new(months.unsigned_abs())),
    }
}

/// Substract the given number of days to the given datetime.
///
/// Returns `None` when it will result in overflow.
pub(crate) fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
    match days.cmp(&0) {
        Ordering::Equal => Some(dt),
        Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
        Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as u64)),
    }
}

#[cfg(test)]
mod tests {

    use chrono::naive::{NaiveDate, NaiveDateTime, NaiveTime};

    use super::*;

    #[test]
    fn test_shift_months() {
        let base = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();

        assert_eq!(
            shift_months(base, 0),
            NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 1),
            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
        );
        assert_eq!(
            shift_months(base, 2),
            NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 3),
            NaiveDate::from_ymd_opt(2020, 4, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, 4),
            NaiveDate::from_ymd_opt(2020, 5, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 5),
            NaiveDate::from_ymd_opt(2020, 6, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, 6),
            NaiveDate::from_ymd_opt(2020, 7, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 7),
            NaiveDate::from_ymd_opt(2020, 8, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 8),
            NaiveDate::from_ymd_opt(2020, 9, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, 9),
            NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 10),
            NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, 11),
            NaiveDate::from_ymd_opt(2020, 12, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 12),
            NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 13),
            NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
        );

        assert_eq!(
            shift_months(base, -1),
            NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -2),
            NaiveDate::from_ymd_opt(2019, 11, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, -3),
            NaiveDate::from_ymd_opt(2019, 10, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -4),
            NaiveDate::from_ymd_opt(2019, 9, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, -5),
            NaiveDate::from_ymd_opt(2019, 8, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -6),
            NaiveDate::from_ymd_opt(2019, 7, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -7),
            NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, -8),
            NaiveDate::from_ymd_opt(2019, 5, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -9),
            NaiveDate::from_ymd_opt(2019, 4, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, -10),
            NaiveDate::from_ymd_opt(2019, 3, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -11),
            NaiveDate::from_ymd_opt(2019, 2, 28).unwrap()
        );
        assert_eq!(
            shift_months(base, -12),
            NaiveDate::from_ymd_opt(2019, 1, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -13),
            NaiveDate::from_ymd_opt(2018, 12, 31).unwrap()
        );

        assert_eq!(
            shift_months(base, 1265),
            NaiveDate::from_ymd_opt(2125, 6, 30).unwrap()
        );
    }

    #[test]
    fn test_shift_months_with_overflow() {
        let base = NaiveDate::from_ymd_opt(2020, 12, 31).unwrap();

        assert_eq!(shift_months(base, 0), base);
        assert_eq!(
            shift_months(base, 1),
            NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 2),
            NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
        );
        assert_eq!(
            shift_months(base, 12),
            NaiveDate::from_ymd_opt(2021, 12, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 18),
            NaiveDate::from_ymd_opt(2022, 6, 30).unwrap()
        );

        assert_eq!(
            shift_months(base, -1),
            NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
        );
        assert_eq!(
            shift_months(base, -2),
            NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -10),
            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
        );
        assert_eq!(
            shift_months(base, -12),
            NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, -18),
            NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
        );
    }

    #[test]
    fn test_shift_months_datetime() {
        let date = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
        let o_clock = NaiveTime::from_hms_opt(1, 2, 3).unwrap();

        let base = NaiveDateTime::new(date, o_clock);

        assert_eq!(
            shift_months(base, 0).date(),
            NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
        );
        assert_eq!(
            shift_months(base, 1).date(),
            NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
        );
        assert_eq!(
            shift_months(base, 2).date(),
            NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
        );
        assert_eq!(shift_months(base, 0).time(), o_clock);
        assert_eq!(shift_months(base, 1).time(), o_clock);
        assert_eq!(shift_months(base, 2).time(), o_clock);
    }
}

Coverage Report

Created: 2025-08-26 07:03

Line	Count	Source
1		// MIT License
2		//
3		// Copyright (c) 2020-2022 Oliver Margetts
4		//
5		// Permission is hereby granted, free of charge, to any person obtaining a copy
6		// of this software and associated documentation files (the "Software"), to deal
7		// in the Software without restriction, including without limitation the rights
8		// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9		// copies of the Software, and to permit persons to whom the Software is
10		// furnished to do so, subject to the following conditions:
11		//
12		// The above copyright notice and this permission notice shall be included in all
13		// copies or substantial portions of the Software.
14		//
15		// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16		// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17		// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18		// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19		// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20		// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21		// SOFTWARE.
22
23		// Copied from chronoutil crate
24
25		//! Contains utility functions for shifting Date objects.
26		use chrono::{DateTime, Datelike, Days, Months, TimeZone};
27		use std::cmp::Ordering;
28
29		/// Shift a date by the given number of months.
30	0	pub(crate) fn shift_months<D>(date: D, months: i32) -> D
31	0	where
32	0	D: Datelike + std::ops::Add<Months, Output = D> + std::ops::Sub<Months, Output = D>,
33		{
34	0	match months.cmp(&0) {
35	0	Ordering::Equal => date,
36	0	Ordering::Greater => date + Months::new(months as u32),
37	0	Ordering::Less => date - Months::new(months.unsigned_abs()),
38		}
39	0	}
40
41		/// Add the given number of months to the given datetime.
42		///
43		/// Returns `None` when it will result in overflow.
44	0	pub(crate) fn add_months_datetime<Tz: TimeZone>(
45	0	dt: DateTime<Tz>,
46	0	months: i32,
47	0	) -> Option<DateTime<Tz>> {
48	0	match months.cmp(&0) {
49	0	Ordering::Equal => Some(dt),
50	0	Ordering::Greater => dt.checked_add_months(Months::new(months as u32)),
51	0	Ordering::Less => dt.checked_sub_months(Months::new(months.unsigned_abs())),
52		}
53	0	}
54
55		/// Add the given number of days to the given datetime.
56		///
57		/// Returns `None` when it will result in overflow.
58	0	pub(crate) fn add_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
59	0	match days.cmp(&0) {
60	0	Ordering::Equal => Some(dt),
61	0	Ordering::Greater => dt.checked_add_days(Days::new(days as u64)),
62	0	Ordering::Less => dt.checked_sub_days(Days::new(days.unsigned_abs() as u64)),
63		}
64	0	}
65
66		/// Substract the given number of months to the given datetime.
67		///
68		/// Returns `None` when it will result in overflow.
69	0	pub(crate) fn sub_months_datetime<Tz: TimeZone>(
70	0	dt: DateTime<Tz>,
71	0	months: i32,
72	0	) -> Option<DateTime<Tz>> {
73	0	match months.cmp(&0) {
74	0	Ordering::Equal => Some(dt),
75	0	Ordering::Greater => dt.checked_sub_months(Months::new(months as u32)),
76	0	Ordering::Less => dt.checked_add_months(Months::new(months.unsigned_abs())),
77		}
78	0	}
79
80		/// Substract the given number of days to the given datetime.
81		///
82		/// Returns `None` when it will result in overflow.
83	0	pub(crate) fn sub_days_datetime<Tz: TimeZone>(dt: DateTime<Tz>, days: i32) -> Option<DateTime<Tz>> {
84	0	match days.cmp(&0) {
85	0	Ordering::Equal => Some(dt),
86	0	Ordering::Greater => dt.checked_sub_days(Days::new(days as u64)),
87	0	Ordering::Less => dt.checked_add_days(Days::new(days.unsigned_abs() as u64)),
88		}
89	0	}
90
91		#[cfg(test)]
92		mod tests {
93
94		use chrono::naive::{NaiveDate, NaiveDateTime, NaiveTime};
95
96		use super::*;
97
98		#[test]
99		fn test_shift_months() {
100		let base = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
101
102		assert_eq!(
103		shift_months(base, 0),
104		NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
105		);
106		assert_eq!(
107		shift_months(base, 1),
108		NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
109		);
110		assert_eq!(
111		shift_months(base, 2),
112		NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
113		);
114		assert_eq!(
115		shift_months(base, 3),
116		NaiveDate::from_ymd_opt(2020, 4, 30).unwrap()
117		);
118		assert_eq!(
119		shift_months(base, 4),
120		NaiveDate::from_ymd_opt(2020, 5, 31).unwrap()
121		);
122		assert_eq!(
123		shift_months(base, 5),
124		NaiveDate::from_ymd_opt(2020, 6, 30).unwrap()
125		);
126		assert_eq!(
127		shift_months(base, 6),
128		NaiveDate::from_ymd_opt(2020, 7, 31).unwrap()
129		);
130		assert_eq!(
131		shift_months(base, 7),
132		NaiveDate::from_ymd_opt(2020, 8, 31).unwrap()
133		);
134		assert_eq!(
135		shift_months(base, 8),
136		NaiveDate::from_ymd_opt(2020, 9, 30).unwrap()
137		);
138		assert_eq!(
139		shift_months(base, 9),
140		NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
141		);
142		assert_eq!(
143		shift_months(base, 10),
144		NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
145		);
146		assert_eq!(
147		shift_months(base, 11),
148		NaiveDate::from_ymd_opt(2020, 12, 31).unwrap()
149		);
150		assert_eq!(
151		shift_months(base, 12),
152		NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
153		);
154		assert_eq!(
155		shift_months(base, 13),
156		NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
157		);
158
159		assert_eq!(
160		shift_months(base, -1),
161		NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
162		);
163		assert_eq!(
164		shift_months(base, -2),
165		NaiveDate::from_ymd_opt(2019, 11, 30).unwrap()
166		);
167		assert_eq!(
168		shift_months(base, -3),
169		NaiveDate::from_ymd_opt(2019, 10, 31).unwrap()
170		);
171		assert_eq!(
172		shift_months(base, -4),
173		NaiveDate::from_ymd_opt(2019, 9, 30).unwrap()
174		);
175		assert_eq!(
176		shift_months(base, -5),
177		NaiveDate::from_ymd_opt(2019, 8, 31).unwrap()
178		);
179		assert_eq!(
180		shift_months(base, -6),
181		NaiveDate::from_ymd_opt(2019, 7, 31).unwrap()
182		);
183		assert_eq!(
184		shift_months(base, -7),
185		NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
186		);
187		assert_eq!(
188		shift_months(base, -8),
189		NaiveDate::from_ymd_opt(2019, 5, 31).unwrap()
190		);
191		assert_eq!(
192		shift_months(base, -9),
193		NaiveDate::from_ymd_opt(2019, 4, 30).unwrap()
194		);
195		assert_eq!(
196		shift_months(base, -10),
197		NaiveDate::from_ymd_opt(2019, 3, 31).unwrap()
198		);
199		assert_eq!(
200		shift_months(base, -11),
201		NaiveDate::from_ymd_opt(2019, 2, 28).unwrap()
202		);
203		assert_eq!(
204		shift_months(base, -12),
205		NaiveDate::from_ymd_opt(2019, 1, 31).unwrap()
206		);
207		assert_eq!(
208		shift_months(base, -13),
209		NaiveDate::from_ymd_opt(2018, 12, 31).unwrap()
210		);
211
212		assert_eq!(
213		shift_months(base, 1265),
214		NaiveDate::from_ymd_opt(2125, 6, 30).unwrap()
215		);
216		}
217
218		#[test]
219		fn test_shift_months_with_overflow() {
220		let base = NaiveDate::from_ymd_opt(2020, 12, 31).unwrap();
221
222		assert_eq!(shift_months(base, 0), base);
223		assert_eq!(
224		shift_months(base, 1),
225		NaiveDate::from_ymd_opt(2021, 1, 31).unwrap()
226		);
227		assert_eq!(
228		shift_months(base, 2),
229		NaiveDate::from_ymd_opt(2021, 2, 28).unwrap()
230		);
231		assert_eq!(
232		shift_months(base, 12),
233		NaiveDate::from_ymd_opt(2021, 12, 31).unwrap()
234		);
235		assert_eq!(
236		shift_months(base, 18),
237		NaiveDate::from_ymd_opt(2022, 6, 30).unwrap()
238		);
239
240		assert_eq!(
241		shift_months(base, -1),
242		NaiveDate::from_ymd_opt(2020, 11, 30).unwrap()
243		);
244		assert_eq!(
245		shift_months(base, -2),
246		NaiveDate::from_ymd_opt(2020, 10, 31).unwrap()
247		);
248		assert_eq!(
249		shift_months(base, -10),
250		NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
251		);
252		assert_eq!(
253		shift_months(base, -12),
254		NaiveDate::from_ymd_opt(2019, 12, 31).unwrap()
255		);
256		assert_eq!(
257		shift_months(base, -18),
258		NaiveDate::from_ymd_opt(2019, 6, 30).unwrap()
259		);
260		}
261
262		#[test]
263		fn test_shift_months_datetime() {
264		let date = NaiveDate::from_ymd_opt(2020, 1, 31).unwrap();
265		let o_clock = NaiveTime::from_hms_opt(1, 2, 3).unwrap();
266
267		let base = NaiveDateTime::new(date, o_clock);
268
269		assert_eq!(
270		shift_months(base, 0).date(),
271		NaiveDate::from_ymd_opt(2020, 1, 31).unwrap()
272		);
273		assert_eq!(
274		shift_months(base, 1).date(),
275		NaiveDate::from_ymd_opt(2020, 2, 29).unwrap()
276		);
277		assert_eq!(
278		shift_months(base, 2).date(),
279		NaiveDate::from_ymd_opt(2020, 3, 31).unwrap()
280		);
281		assert_eq!(shift_months(base, 0).time(), o_clock);
282		assert_eq!(shift_months(base, 1).time(), o_clock);
283		assert_eq!(shift_months(base, 2).time(), o_clock);
284		}
285		}