From 1fbc26374e2647dbf09ff2c2736d6ef6192106b3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Fl=C3=BCgge?= <seflue@web.de>
Date: Wed, 5 Jun 2024 19:52:42 +0200
Subject: [PATCH] Fix date arithmetic issues with month (#748)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* test(date): expose bug with test

* fix(date): handle end of month correctly

- handle 30d month after 31d month correctly
- handle February correctly including leap years

---------

Co-authored-by: Sebastian Flügge <seflue@users.noreply.github.com>
---
 lua/orgmode/objects/date.lua       | 26 +++++++++++++-
 tests/plenary/object/date_spec.lua | 56 ++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)

diff --git a/lua/orgmode/objects/date.lua b/lua/orgmode/objects/date.lua
index 96f3a04a7..0290d3778 100644
--- a/lua/orgmode/objects/date.lua
+++ b/lua/orgmode/objects/date.lua
@@ -481,7 +481,8 @@ function Date:end_of(span)
   end
 
   if span == 'month' then
-    return self:add({ month = 1 }):start_of('month'):adjust('-1d'):end_of('day')
+    local date = os.date('*t', self.timestamp)
+    return self:set({ day = Date._days_of_month(date) }):end_of('day')
   end
 
   return self
@@ -517,6 +518,7 @@ end
 ---@return OrgDate
 function Date:add(opts)
   opts = opts or {}
+  ---@type table
   local date = os.date('*t', self.timestamp)
   for opt, val in pairs(opts) do
     if opt == 'week' then
@@ -525,9 +527,31 @@ function Date:add(opts)
     end
     date[opt] = date[opt] + val
   end
+  if opts['month'] then
+    date['day'] = math.min(date['day'], Date._days_of_month(date))
+  end
   return self:from_time_table(date)
 end
 
+---@param date table
+---@return number
+function Date._days_of_month(date)
+  local days_of = { 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31 }
+  if date.month == 2 then
+    return Date._days_of_february(date.year)
+  else
+    return days_of[date.month]
+  end
+end
+
+function Date._days_of_february(year)
+  return Date._is_leap_year(year) and 29 or 28
+end
+
+function Date._is_leap_year(year)
+  return year % 400 == 0 or (year % 100 ~= 0 and year % 4 == 0)
+end
+
 ---@param opts table
 ---@return OrgDate
 function Date:subtract(opts)
diff --git a/tests/plenary/object/date_spec.lua b/tests/plenary/object/date_spec.lua
index d5c1f3249..eec71f290 100644
--- a/tests/plenary/object/date_spec.lua
+++ b/tests/plenary/object/date_spec.lua
@@ -303,6 +303,14 @@ describe('Date object', function()
     assert.are.same('2021-05-31 Mon 23:59', date:to_string())
   end)
 
+  it('should properly handle end of month', function()
+    local date = Date.from_string('2021-05-12')
+    date = date:end_of('month')
+    assert.are.same('2021-05-31 Mon', date:to_string())
+    date = date:end_of('month')
+    assert.are.same('2021-05-31 Mon', date:to_string())
+  end)
+
   it('should add/subtract/set date', function()
     local date = Date.from_string('2021-05-12 14:00')
     date = date:add({ week = 2 })
@@ -810,4 +818,52 @@ describe('Date object', function()
     local end_of_2021 = Date.from_string('2021-12-31')
     assert.are.same('52', end_of_2021:get_week_number())
   end)
+
+  it('should add month correctly | long month + short month', function()
+    local date = Date.from_string('2021-05-31')
+    assert.are.same('2021-05-31 Mon', date:to_string())
+    assert.are.same('2021-06-30 Wed', date:add({ month = 1 }):to_string())
+  end)
+
+  it('should add month correctly | short month + long month', function()
+    local date = Date.from_string('2021-04-30')
+    assert.are.same('2021-04-30 Fri', date:to_string())
+    assert.are.same('2021-05-30 Sun', date:add({ month = 1 }):to_string())
+  end)
+
+  it('should add month correctly | long month + february', function()
+    local date = Date.from_string('2021-01-31')
+    assert.are.same('2021-01-31 Sun', date:to_string())
+    assert.are.same('2021-02-28 Sun', date:add({ month = 1 }):to_string())
+  end)
+
+  it('should add month correctly | long month + february in leap year', function()
+    local date = Date.from_string('2024-01-31')
+    assert.are.same('2024-01-31 Wed', date:to_string())
+    assert.are.same('2024-02-29 Thu', date:add({ month = 1 }):to_string())
+  end)
+
+  it('should calculate end of month correctly | long month', function()
+    local date = Date.from_string('2021-05-31')
+    assert.are.same('2021-05-31 Mon', date:to_string())
+    assert.are.same('2021-05-31 Mon', date:end_of('month'):to_string())
+  end)
+
+  it('should calculate end of month correctly | short month', function()
+    local date = Date.from_string('2021-04-30')
+    assert.are.same('2021-04-30 Fri', date:to_string())
+    assert.are.same('2021-04-30 Fri', date:end_of('month'):to_string())
+  end)
+
+  it('should calculate end of month correctly | february', function()
+    local date = Date.from_string('2021-02-28')
+    assert.are.same('2021-02-28 Sun', date:to_string())
+    assert.are.same('2021-02-28 Sun', date:end_of('month'):to_string())
+  end)
+
+  it('should calculate end of month correctly | february leap-year', function()
+    local date = Date.from_string('2024-02-29')
+    assert.are.same('2024-02-29 Thu', date:to_string())
+    assert.are.same('2024-02-29 Thu', date:end_of('month'):to_string())
+  end)
 end)