Calendar next/prev supports count prefix (#438)

Author: jgollenz

lua/orgmode/objects/calendar.lua

@@ -32,15 +32,20 @@ function Calendar.new(data)
   return Calendar
 end
 
+local width = 36
+local height = 10
+local x_offset = 1 -- one border cell
+local y_offset = 2 -- one border cell and one padding cell
+
 function Calendar.open()
   local opts = {
     relative = 'editor',
-    width = 36,
-    height = Calendar.clearable and 11 or 10,
+    width = width,
+    height = Calendar.clearable and height + 1 or height,
     style = 'minimal',
     border = 'single',
-    row = vim.o.lines / 2 - 4,
-    col = vim.o.columns / 2 - 20,
+    row = vim.o.lines / 2 - (y_offset + height) / 2,
+    col = vim.o.columns / 2 - (x_offset + width) / 2,
   }
 
   Calendar.buf = vim.api.nvim_create_buf(false, true)
@@ -85,7 +90,7 @@ function Calendar.open()
   if Calendar.date then
     search_day = Calendar.date:format('%d')
   end
-  vim.fn.cursor({ 2, 0 })
+  vim.fn.cursor(2, 0)
   vim.fn.search(search_day, 'W')
   return Promise.new(function(resolve)
     Calendar.callback = resolve
@@ -94,55 +99,69 @@ end
 
 function Calendar.render()
   vim.api.nvim_buf_set_option(Calendar.buf, 'modifiable', true)
+
+  local cal_rows = { {}, {}, {}, {}, {}, {} } -- the calendar rows
   local start_from_sunday = config.calendar_week_start_day == 0
+  local weekday_row = { 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun' }
 
-  local first_row = { 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun' }
   if start_from_sunday then
-    first_row = { 'Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat' }
+    weekday_row = { 'Sun', 'Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat' }
   end
-  local content = { {}, {}, {}, {}, {}, {} }
+
+  -- construct title (Month YYYY)
+  local title = Calendar.month:format('%B %Y')
+  title = string.rep(' ', math.floor((width - title:len()) / 2)) .. title
+
+  -- insert whitespace before first day of month
   local start_weekday = Calendar.month:get_isoweekday()
   if start_from_sunday then
     start_weekday = Calendar.month:get_weekday()
   end
   while start_weekday > 1 do
-    table.insert(content[1], '  ')
+    table.insert(cal_rows[1], '  ')
     start_weekday = start_weekday - 1
   end
-  local today = Date.today()
-  local is_today_month = today:is_same(Calendar.month, 'month')
+
+  -- insert dates into cal_rows
   local dates = Calendar.month:get_range_until(Calendar.month:end_of('month'))
-  local month = Calendar.month:format('%B %Y')
-  month = string.rep(' ', math.floor((36 - month:len()) / 2)) .. month
-  local start_row = 1
+  local current_row = 1
   for _, date in ipairs(dates) do
-    table.insert(content[start_row], date:format('%d'))
-    if #content[start_row] % 7 == 0 then
-      start_row = start_row + 1
+    table.insert(cal_rows[current_row], date:format('%d'))
+    if #cal_rows[current_row] % 7 == 0 then
+      current_row = current_row + 1
     end
   end
-  local value = vim.tbl_map(function(item)
+
+  -- add spacing between the calendar cells
+  local content = vim.tbl_map(function(item)
     return ' ' .. table.concat(item, '   ') .. ' '
-  end, content)
-  first_row = ' ' .. table.concat(first_row, '  ')
-  table.insert(value, 1, first_row)
-  table.insert(value, 1, month)
-  table.insert(value, ' [<] - prev month  [>] - next month')
-  table.insert(value, ' [.] - today   [Enter] - select day')
+  end, cal_rows)
+  weekday_row = ' ' .. table.concat(weekday_row, '  ')
+
+  -- put it all together
+  table.insert(content, 1, weekday_row)
+  table.insert(content, 1, title)
+  -- TODO: redundant, since it's static data
+  table.insert(content, ' [<] - prev month  [>] - next month')
+  table.insert(content, ' [.] - today   [Enter] - select day')
   if Calendar.clearable then
-    table.insert(value, ' [r] Clear date')
+    table.insert(content, ' [r] Clear date')
   end
 
-  vim.api.nvim_buf_set_lines(Calendar.buf, 0, -1, true, value)
+  vim.api.nvim_buf_set_lines(Calendar.buf, 0, -1, true, content)
   vim.api.nvim_buf_clear_namespace(Calendar.buf, Calendar.namespace, 0, -1)
   if Calendar.clearable then
-    vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #value - 3, 0, -1)
+    vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #content - 3, 0, -1)
   end
-  vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #value - 2, 0, -1)
-  vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #value - 1, 0, -1)
+  vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #content - 2, 0, -1)
+  vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'Comment', #content - 1, 0, -1)
+
+  -- highlight the cell of the current day
+  local today = Date.today()
+  local is_today_month = today:is_same(Calendar.month, 'month')
   if is_today_month then
     local day_formatted = today:format('%d')
-    for i, line in ipairs(value) do
+    for i, line in ipairs(content) do
       local from, to = line:find('%s' .. day_formatted .. '%s')
       if from and to then
         vim.api.nvim_buf_add_highlight(Calendar.buf, Calendar.namespace, 'OrgCalendarToday', i - 1, from - 1, to)
@@ -154,14 +173,14 @@ function Calendar.render()
 end
 
 function Calendar.forward()
-  Calendar.month = Calendar.month:add({ month = 1 })
+  Calendar.month = Calendar.month:add({ month = vim.v.count1 })
   Calendar.render()
-  vim.fn.cursor({ 2, 0 })
+  vim.fn.cursor(2, 0)
   vim.fn.search('01')
 end
 
 function Calendar.backward()
-  Calendar.month = Calendar.month:subtract({ month = 1 })
+  Calendar.month = Calendar.month:subtract({ month = vim.v.count1 })
   Calendar.render()
   vim.fn.cursor('$', 0)
   vim.fn.search([[\d\d]], 'b')
@@ -173,7 +192,7 @@ function Calendar.cursor_right()
     local curr_line = vim.fn.getline('.')
     local offset = curr_line:sub(col + 1, #curr_line):find('%d%d')
     if offset ~= nil then
-      vim.fn.cursor({ line, col + offset })
+      vim.fn.cursor(line, col + offset)
     end
   end
 end