From 2a10172bf71f59634f79b97eea1200266b1c603b Mon Sep 17 00:00:00 2001
From: Kristijan Husak <husakkristijan@gmail.com>
Date: Mon, 23 Sep 2024 09:23:14 +0200
Subject: [PATCH] feat(dev): Add persist option to load_file method

---
 lua/orgmode/files/init.lua  | 29 ++++++++++++++++++++++++++---
 tests/plenary/init_spec.lua | 23 ++++++++++++++++++++++-
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/lua/orgmode/files/init.lua b/lua/orgmode/files/init.lua
index 0b200897..aeb65fc9 100644
--- a/lua/orgmode/files/init.lua
+++ b/lua/orgmode/files/init.lua
@@ -8,6 +8,9 @@ local Listitem = require('orgmode.files.elements.listitem')
 ---@class OrgFilesOpts
 ---@field paths string | string[]
 
+---@class OrgLoadFileOpts
+---@field persist boolean Persist the file in the list of loaded files if it belongs to path
+
 ---@class OrgFiles
 ---@field paths string[]
 ---@field files table<string, OrgFile> table with files that are part of paths
@@ -52,6 +55,7 @@ function OrgFiles:load(force)
   end)
 end
 
+---@deprecated Use `load_file` with `persist` option instead
 ---@param filename string
 ---@return OrgPromise<OrgFile | false>
 function OrgFiles:add_to_paths(filename)
@@ -153,25 +157,44 @@ function OrgFiles:filenames()
   end, self:all())
 end
 
+---@param filename string
+---@param opts? OrgLoadFileOpts
 ---@return OrgPromise<OrgFile | false>
-function OrgFiles:load_file(filename)
+function OrgFiles:load_file(filename, opts)
+  opts = opts or {}
   filename = vim.fn.resolve(vim.fn.fnamemodify(filename, ':p'))
+
+  local persist_if_required = function(file)
+    ---@cast file OrgFile
+    if self.files[filename] or not opts.persist then
+      return
+    end
+    local all_paths = self:_files()
+    if vim.tbl_contains(all_paths, filename) then
+      self.files[filename] = file
+    end
+  end
+
   local file = self.all_files[filename]
   if file then
+    persist_if_required(file)
     return file:reload()
   end
 
   return OrgFile.load(filename):next(function(orgfile)
     if orgfile then
+      persist_if_required(file)
       self.all_files[filename] = orgfile
     end
     return orgfile
   end)
 end
 
+---@param filename string
+---@param opts? OrgLoadFileOpts
 ---@return OrgFile | false
-function OrgFiles:load_file_sync(filename, timeout)
-  return self:load_file(filename):wait(timeout)
+function OrgFiles:load_file_sync(filename, opts, timeout)
+  return self:load_file(filename, opts):wait(timeout)
 end
 
 ---@param filename string
diff --git a/tests/plenary/init_spec.lua b/tests/plenary/init_spec.lua
index c4586876..71855096 100644
--- a/tests/plenary/init_spec.lua
+++ b/tests/plenary/init_spec.lua
@@ -1,5 +1,4 @@
 local orgmode = require('orgmode')
-local Date = require('orgmode.objects.date')
 
 describe('Init', function()
   local org = orgmode.setup({
@@ -37,6 +36,28 @@ describe('Init', function()
     assert.are.same({ 'NESTED', 'OFFICE', 'PRIVATE', 'PROJECT', 'WORK' }, org.files:get_tags())
   end)
 
+  it('should load file and persist to files if it belongs to path', function()
+    local fname = vim.fn.resolve(vim.fn.tempname() .. '.org')
+    vim.fn.writefile({ '* Appended' }, fname)
+
+    assert.is.Nil(org.files.files[fname])
+    assert.are.same({}, org.files:find_headlines_by_title('Appended'))
+    assert.are.same({ vim.fn.getcwd() .. '/tests/plenary/fixtures/*' }, org.files.paths)
+
+    -- Not added because it does not belong to defined path
+    org.files:load_file_sync(fname, { persist = true })
+    assert.is.Nil(org.files.files[fname])
+
+    org.files.files[todo_file] = nil
+
+    org.files:load_file_sync(todo_file)
+    -- Not added because persist was not provided
+    assert.is.Nil(org.files.files[fname])
+
+    org.files:load_file_sync(todo_file, { persist = true })
+    assert.is.Not.Nil(org.files.files[todo_file])
+  end)
+
   it('should append files to paths', function()
     local fname = vim.fn.resolve(vim.fn.tempname() .. '.org')
     vim.fn.writefile({ '* Appended' }, fname)