feat(capture): Add support for template datetree (#682)

Closes #672

Author: kristijanhusak

DOCS.md

@@ -488,6 +488,8 @@ Templates have the following fields:
   * `template` (`string|string[]`) — body of the template that will be used when creating capture
   * `target` (`string?`) — name of the file to which the capture content will be added. If the target is not specified, the content will be added to the [`org_default_notes_file`](#orgdefaultnotesfile) file
   * `headline` (`string?`) — title of the headline after which the capture content will be added. If no headline is specified, the content will be appended to the end of the file
+  * `datetree (boolean | { time_prompt: boolean })` — Create a [date tree](https://orgmode.org/manual/Template-elements.html#FOOT84) with current day in the target file and put the capture content there.
+    When `time_prompt = true`, open up a date picker to select a date before opening up a capture buffer.
   * `properties` (`table?`):
     * `empty_lines` (`table|number?`) — if the value is a number, then empty lines are added before and after the content. If the value is a table, then the following fields are expected:
       * `before` (`integer?`) — add empty lines to the beginning of the content

lua/orgmode/capture/init.lua

@@ -6,6 +6,7 @@ local Menu = require('orgmode.ui.menu')
 local Range = require('orgmode.files.elements.range')
 local CaptureWindow = require('orgmode.capture.window')
 local Date = require('orgmode.objects.date')
+local Datetree = require('orgmode.capture.template.datetree')
 
 ---@class OrgProcessRefileOpts
 ---@field source_headline OrgHeadline
@@ -18,7 +19,7 @@ local Date = require('orgmode.objects.date')
 ---@field template OrgCaptureTemplate
 ---@field source_file OrgFile
 ---@field source_headline? OrgHeadline
----@field destination_file? OrgFile
+---@field destination_file OrgFile
 ---@field destination_headline? OrgHeadline
 
 ---@class OrgCapture
@@ -55,6 +56,15 @@ function Capture:open_template(template)
     end,
   })
 
+  if template:has_input_prompts() then
+    return template:prompt_for_inputs():next(function(proceed)
+      if not proceed then
+        return utils.echo_info('Canceled.')
+      end
+      return self._window:open()
+    end)
+  end
+
   return self._window:open()
 end
 
@@ -124,31 +134,37 @@ end
 function Capture:_refile_from_capture_buffer(opts)
   local target_level = 0
   local target_line = -1
+  local destination_file = opts.destination_file
+  local destination_headline = opts.destination_headline
 
-  if opts.destination_headline then
-    target_level = opts.destination_headline:get_level()
-    target_line = opts.destination_headline:get_range().end_line
+  if opts.template.datetree then
+    destination_headline = Datetree:new({ files = self.files }):create(opts.template)
+  end
+
+  if destination_headline then
+    target_level = destination_headline:get_level()
+    target_line = destination_headline:get_range().end_line
   end
 
   local lines = opts.source_file.lines
 
   if opts.source_headline then
     lines = opts.source_headline:get_lines()
-    if opts.destination_headline or opts.source_headline:get_level() > 1 then
+    if destination_headline or opts.source_headline:get_level() > 1 then
       lines = self:_adapt_headline_level(opts.source_headline, target_level, false)
     end
   end
 
   lines = opts.template:apply_properties_to_lines(lines)
 
   self.files
-    :update_file(opts.destination_file.filename, function(file)
+    :update_file(destination_file.filename, function(file)
       local range = self:_get_destination_range_without_empty_lines(Range.from_line(target_line))
       vim.api.nvim_buf_set_lines(file:bufnr(), range.start_line, range.end_line, false, lines)
     end)
     :wait()
 
-  utils.echo_info(('Wrote %s'):format(opts.destination_file.filename))
+  utils.echo_info(('Wrote %s'):format(destination_file.filename))
   self:kill()
   return true
 end
@@ -413,8 +429,7 @@ end
 ---@private
 ---@return OrgProcessCaptureOpts | false
 function Capture:_get_refile_vars()
-  local target = self._window.template.target
-  local file = vim.fn.resolve(vim.fn.fnamemodify(target, ':p'))
+  local file = self._window.template:get_target()
 
   if vim.fn.filereadable(file) == 0 then
     local choice = vim.fn.confirm(('Refile destination %s does not exist. Create now?'):format(file), '&Yes\n&No')

lua/orgmode/capture/template/datetree.lua

@@ -0,0 +1,163 @@
+local utils = require('orgmode.utils')
+local Date = require('orgmode.objects.date')
+
+---@class OrgDatetree
+---@field files OrgFiles
+local Datetree = {}
+Datetree.__index = Datetree
+
+---@param opts { files: OrgFiles }
+function Datetree:new(opts)
+  return setmetatable({
+    files = opts.files,
+  }, Datetree)
+end
+
+---@param template OrgCaptureTemplate
+---@return OrgHeadline
+function Datetree:create(template)
+  local date = template:get_datetree_date()
+  local destination_file = self.files:get(template:get_target())
+  local result = self:_get_datetree_destination(destination_file, date)
+
+  if result.create then
+    self.files
+      :update_file(destination_file.filename, function(file)
+        vim.api.nvim_buf_set_lines(file:bufnr(), result.target_line, result.target_line, false, result.content)
+      end)
+      :wait()
+    destination_file = destination_file:reload_sync()
+  end
+
+  return destination_file:get_closest_headline({ result.headline_at, 0 })
+end
+
+---@private
+---@param destination_file OrgFile
+---@param date OrgDate
+---@return { create: boolean, target_line: number, content: string[], headline_at: number }
+function Datetree:_get_datetree_destination(destination_file, date)
+  local year_date = date:format('%Y')
+  local month_date = date:start_of('month')
+  local month_date_str = date:format('%Y-%m %B')
+  local day_date = date:format('%Y-%m-%d %A')
+  local year_headline = utils.find(destination_file:get_top_level_headlines(), function(headline)
+    return headline:get_title() == year_date
+  end)
+
+  if not year_headline then
+    local target_line = self:_get_insert_year_at(destination_file, year_date)
+    return {
+      create = true,
+      target_line = target_line,
+      headline_at = target_line + 3,
+      content = {
+        '* ' .. year_date,
+        '** ' .. month_date_str,
+        '*** ' .. day_date,
+      },
+    }
+  end
+
+  local month_headline = utils.find(year_headline:get_child_headlines(), function(month)
+    return month:get_title() == month_date_str
+  end)
+
+  if not month_headline then
+    local target_line = self:_get_insert_month_at(year_headline, month_date)
+    return {
+      create = true,
+      target_line = target_line,
+      headline_at = target_line + 2,
+      content = {
+        '** ' .. month_date_str,
+        '*** ' .. day_date,
+      },
+    }
+  end
+
+  local month_headlines = month_headline:get_child_headlines()
+  local day_headline = utils.find(month_headlines, function(day)
+    return day:get_title() == day_date
+  end)
+
+  if not day_headline then
+    local target_line = self:_get_insert_day_at(month_headline, date)
+
+    return {
+      create = true,
+      target_line = target_line,
+      headline_at = target_line + 1,
+      content = {
+        '*** ' .. day_date,
+      },
+    }
+  end
+
+  return {
+    create = false,
+    headline_at = day_headline:get_range().start_line,
+  }
+end
+
+---@private
+---@param destination_file OrgFile
+---@param year_date string -- year in format YYYY
+---@return number
+function Datetree:_get_insert_year_at(destination_file, year_date)
+  local future_year_headline = utils.find(destination_file:get_top_level_headlines(), function(headline)
+    local get_year = headline:get_title():match('^%d%d%d%d$')
+    return get_year and tonumber(get_year) > tonumber(year_date)
+  end)
+
+  if future_year_headline then
+    return future_year_headline:get_range().start_line - 1
+  end
+
+  return #destination_file.lines
+end
+
+---@private
+---@param year_headline OrgHeadline
+---@param month_date OrgDate
+---@return number
+function Datetree:_get_insert_month_at(year_headline, month_date)
+  local future_month_headline = utils.find(year_headline:get_child_headlines(), function(headline)
+    local year_num, month_num = headline:get_title():match('^(%d%d%d%d)%-(%d%d)%s+%w+$')
+    if year_num and month_num then
+      local timestamp = os.time({ year = year_num, month = month_num, day = 1 })
+      return timestamp > month_date.timestamp
+    end
+    return false
+  end)
+
+  if future_month_headline then
+    return future_month_headline:get_range().start_line - 1
+  end
+
+  return year_headline:get_range().end_line
+end
+
+---@private
+---@param month_headline OrgHeadline
+---@param date OrgDate
+---@return number
+function Datetree:_get_insert_day_at(month_headline, date)
+  local future_day_headline = utils.find(month_headline:get_child_headlines(), function(headline)
+    local year_num, month_num, day_num = headline:get_title():match('^(%d%d%d%d)%-(%d%d)%-(%d%d)%s+%w+$')
+    return year_num
+      and Date.from_table({
+        year = year_num,
+        month = month_num,
+        day = day_num,
+      }):is_after(date, 'day')
+  end)
+
+  if future_day_headline then
+    return future_day_headline:get_range().start_line - 1
+  end
+
+  return month_headline:get_range().end_line
+end
+
+return Datetree

lua/orgmode/capture/template/init.lua

@@ -1,6 +1,8 @@
 local TemplateProperties = require('orgmode.capture.template.template_properties')
 local Date = require('orgmode.objects.date')
 local utils = require('orgmode.utils')
+local Calendar = require('orgmode.objects.calendar')
+local Promise = require('orgmode.utils.promise')
 
 local expansions = {
   ['%f'] = function()
@@ -32,15 +34,20 @@ local expansions = {
   end,
 }
 
+---@alias OrgCaptureTemplateDatetree boolean | { time_prompt: boolean, date?: OrgDate }
+
 ---@class OrgCaptureTemplate
----@field description string
----@field template string|string[]
----@field target string?
----@field headline string?
----@field properties OrgCaptureTemplateProperties
----@field subtemplates table<string, OrgCaptureTemplate>
+---@field description? string
+---@field template? string|string[]
+---@field target? string
+---@field datetree? OrgCaptureTemplateDatetree
+---@field headline? string
+---@field properties? OrgCaptureTemplateProperties
+---@field subtemplates? table<string, OrgCaptureTemplate>
 local Template = {}
 
+---@param opts OrgCaptureTemplate
+---@return OrgCaptureTemplate
 function Template:new(opts)
   opts = opts or {}
 
@@ -51,6 +58,7 @@ function Template:new(opts)
     headline = { opts.headline, 'string', true },
     properties = { opts.properties, 'table', true },
     subtemplates = { opts.subtemplates, 'table', true },
+    datetree = { opts.datetree, { 'boolean', 'table' }, true },
   })
 
   local this = {}
@@ -59,6 +67,7 @@ function Template:new(opts)
   this.target = self:_compile(opts.target or '')
   this.headline = opts.headline
   this.properties = TemplateProperties:new(opts.properties)
+  this.datetree = opts.datetree
 
   this.subtemplates = {}
   for key, subtemplate in pairs(opts.subtemplates or {}) do
@@ -89,10 +98,39 @@ function Template:compile()
   if type(content) == 'table' then
     content = table.concat(content, '\n')
   end
-  content = self:_compile(content)
+  content = self:_compile(content or '')
   return vim.split(content, '\n', { plain = true })
 end
 
+function Template:has_input_prompts()
+  return self.datetree and type(self.datetree) == 'table' and self.datetree.time_prompt
+end
+
+function Template:prompt_for_inputs()
+  if not self:has_input_prompts() then
+    return Promise.resolve(true)
+  end
+  return Calendar.new({ date = Date.now() }):open():next(function(date)
+    if date then
+      self.datetree.date = date
+      return true
+    end
+    return false
+  end)
+end
+
+function Template:get_datetree_date()
+  if self:has_input_prompts() then
+    return self.datetree.date
+  end
+  return Date.today()
+end
+
+---@return string
+function Template:get_target()
+  return vim.fn.resolve(vim.fn.fnamemodify(self.target, ':p'))
+end
+
 ---@param lines string[]
 ---@return string[]
 function Template:apply_properties_to_lines(lines)

lua/orgmode/files/file.lua

@@ -170,6 +170,18 @@ function OrgFile:get_headlines()
   end, matches)
 end
 
+memoize('get_top_level_headlines')
+---@return OrgHeadline[]
+function OrgFile:get_top_level_headlines()
+  if self:is_archive_file() then
+    return {}
+  end
+  local matches = self:get_ts_matches('(document (section (headline) @headline))')
+  return vim.tbl_map(function(match)
+    return Headline:new(match.headline.node, self)
+  end, matches)
+end
+
 memoize('get_headlines_including_archived')
 ---@return OrgHeadline[]
 function OrgFile:get_headlines_including_archived()
@@ -181,8 +193,9 @@ end
 
 ---@param title string
 ---@param exact? boolean
+---@param search_from_end? boolean
 ---@return OrgHeadline[]
-function OrgFile:find_headlines_by_title(title, exact)
+function OrgFile:find_headlines_by_title(title, exact, search_from_end)
   return vim.tbl_filter(function(item)
     local pattern = '^' .. vim.pesc(title:lower())
     if exact then
@@ -195,7 +208,9 @@ end
 ---@param title string
 ---@return OrgHeadline | nil
 function OrgFile:find_headline_by_title(title)
-  return self:find_headlines_by_title(title, true)[1]
+  return utils.find(self:get_headlines(), function(item)
+    return item:get_title():lower() == title:lower()
+  end)
 end
 
 ---@return OrgHeadline[]