From 0a8be4a8b93c065fa8c27e706ff8b69654437b22 Mon Sep 17 00:00:00 2001
From: Michael Ball <michael@mball.co>
Date: Sun, 19 May 2024 15:56:51 -0700
Subject: [PATCH] Add support for annotating lua model files

---
 lapis/cmd/actions/annotate.lua  | 14 +++++++++++---
 lapis/cmd/actions/annotate.moon | 10 +++++++---
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/lapis/cmd/actions/annotate.lua b/lapis/cmd/actions/annotate.lua
index 177a7d2..06aa740 100644
--- a/lapis/cmd/actions/annotate.lua
+++ b/lapis/cmd/actions/annotate.lua
@@ -23,7 +23,7 @@ build_command = function(cmd, config)
   do
     local password = config.postgres.password
     if password then
-      table.insert(1, command, "PGPASSWORD=" .. tostring(shell_escape(password)))
+      table.insert(command, 1, "PGPASSWORD=" .. tostring(shell_escape(password)))
     end
   end
   do
@@ -229,8 +229,16 @@ annotate_model = function(config, fname, options)
   end
   local header = table.concat(header_lines, "\n") .. "\n"
   local updated_source = replace_header(source, header)
-  if not (updated_source) then
-    updated_source = source:gsub("class ", tostring(header) .. "class ", 1)
+  if fname:match(".lua$") and not updated_source then
+    local lua_declaration = "local %w+ = Model:extend"
+    if not source:match(lua_declaration) then
+      print("\tLine matching '" .. tostring(lua_declaration) .. "' not found")
+    end
+    updated_source = source:gsub(lua_declaration, tostring(header) .. "%1 ", 1)
+  else
+    if not updated_source then
+      updated_source = source:gsub("class ", tostring(header) .. "class ", 1)
+    end
   end
   local source_out = assert(io.open(fname, "w"))
   source_out:write(updated_source)
diff --git a/lapis/cmd/actions/annotate.moon b/lapis/cmd/actions/annotate.moon
index 674f748..e3ec921 100644
--- a/lapis/cmd/actions/annotate.moon
+++ b/lapis/cmd/actions/annotate.moon
@@ -17,7 +17,7 @@ build_command = (cmd, config) ->
   command = { cmd }
 
   if password = config.postgres.password
-    table.insert 1, command, "PGPASSWORD=#{shell_escape password}"
+    table.insert command, 1, "PGPASSWORD=#{shell_escape password}"
 
   if host = config.postgres.host
     table.insert command, "-h #{shell_escape host}"
@@ -149,8 +149,12 @@ annotate_model = (config, fname, options={}) ->
 
   updated_source = replace_header source, header
 
-  -- TODO: this is kinda sloppy and only works with MoonScript
-  unless updated_source
+  -- TODO: this is kinda sloppy
+  if fname\match(".lua$") and not updated_source
+    lua_declaration = "local %w+ = Model:extend"
+    print "\tLine matching '#{lua_declaration}' not found" if not source\match lua_declaration
+    updated_source = source\gsub lua_declaration, "#{header}%1 ", 1
+  else if not updated_source
     updated_source = source\gsub "class ", "#{header}class ", 1
 
   source_out = assert io.open fname, "w"