Merge pull request #580 from lethosor/lua-enable

Allow Lua scripts to implement enable/disable commands and specify other behavior (e.g. "module mode") that they implement

Lua API.rst

@@ -3457,6 +3457,30 @@ Note that this function lets errors propagate to the caller.
   This will not be an issue in most cases.
   This function also permits circular dependencies of scripts.
 
+* ``dfhack.reqscript(name)`` or ``reqscript(name)``
+
+  Nearly identical to script_environment() but requires scripts being loaded to
+  include a line similar to::
+
+      --@ module = true
+
+  This is intended to only allow scripts that take appropriate action when used
+  as a module to be loaded.
+
+Enabling and disabling scripts
+==============================
+
+Scripts can choose to recognize the built-in ``enable`` and ``disable`` commands
+by including the following line anywhere in their file::
+
+    --@ enable = true
+
+When the ``enable`` and ``disable`` commands are invoked, a ``dfhack_flags``
+table will be passed to the script with the following fields set:
+
+* ``enable``: Always true if the script is being enabled *or* disabled
+* ``enable_state``: True if the script is being enabled, false otherwise
+
 Save init script
 ================
 

NEWS

@@ -4,6 +4,8 @@ DFHack Future
         Developer plugins can be ignored on startup by setting the DFHACK_NO_DEV_PLUGINS environment variable
         The console on Linux and OS X now recognizes keyboard input between prompts
     Lua
+        Scripts can be enabled with the built-in enable/disable commands
+        A new function, reqscript(), is available as a safer alternative to script_environment()
     New internal commands
     New plugins
     New scripts

library/Core.cpp

@@ -288,6 +288,10 @@ namespace {
         const string *pcmd;
         vector<string> *pargs;
     };
+    struct ScriptEnableState {
+        const string *pcmd;
+        bool pstate;
+    };
 }
 
 static bool init_run_script(color_ostream &out, lua_State *state, void *info)
@@ -315,6 +319,30 @@ static command_result runLuaScript(color_ostream &out, std::string name, vector<
     return ok ? CR_OK : CR_FAILURE;
 }
 
+static bool init_enable_script(color_ostream &out, lua_State *state, void *info)
+{
+    auto args = (ScriptEnableState*)info;
+    if (!lua_checkstack(state, 4))
+        return false;
+    Lua::PushDFHack(state);
+    lua_getfield(state, -1, "enable_script");
+    lua_remove(state, -2);
+    lua_pushstring(state, args->pcmd->c_str());
+    lua_pushboolean(state, args->pstate);
+    return true;
+}
+
+static command_result enableLuaScript(color_ostream &out, std::string name, bool state)
+{
+    ScriptEnableState data;
+    data.pcmd = &name;
+    data.pstate = state;
+
+    bool ok = Lua::RunCoreQueryLoop(out, Lua::Core::State, init_enable_script, &data);
+
+    return ok ? CR_OK : CR_FAILURE;
+}
+
 static command_result runRubyScript(color_ostream &out, PluginManager *plug_mgr, std::string name, vector<string> &args)
 {
     if (!plug_mgr->ruby || !plug_mgr->ruby->is_enabled())
@@ -633,24 +661,43 @@ command_result Core::runCommand(color_ostream &con, const std::string &first_, v
 
                 for (size_t i = 0; i < parts.size(); i++)
                 {
-                    Plugin * plug = plug_mgr->getPluginByName(parts[i]);
+                    std::string part = parts[i];
+                    if (part.find('\\') != std::string::npos)
+                    {
+                        con.printerr("Replacing backslashes with forward slashes in \"%s\"\n", part.c_str());
+                        for (size_t j = 0; j < part.size(); j++)
+                        {
+                            if (part[j] == '\\')
+                                part[j] = '/';
+                        }
+                    }
+
+                    Plugin * plug = plug_mgr->getPluginByName(part);
 
                     if(!plug)
+                    {
+                        std::string lua = findScript(this->p->getPath(), part + ".lua");
+                        if (lua.size())
+                        {
+                            res = enableLuaScript(con, part, enable);
+                        }
+                        else
                         {
                             res = CR_NOT_FOUND;
-                        con.printerr("No such plugin: %s\n", parts[i].c_str());
+                            con.printerr("No such plugin or Lua script: %s\n", part.c_str());
+                        }
                     }
                     else if (!plug->can_set_enabled())
                     {
                         res = CR_NOT_IMPLEMENTED;
-                        con.printerr("Cannot %s plugin: %s\n", first.c_str(), parts[i].c_str());
+                        con.printerr("Cannot %s plugin: %s\n", first.c_str(), part.c_str());
                     }
                     else
                     {
                         res = plug->set_enabled(con, enable);
 
                         if (res != CR_OK || plug->is_enabled() != enable)
-                            con.printerr("Could not %s plugin: %s\n", first.c_str(), parts[i].c_str());
+                            con.printerr("Could not %s plugin: %s\n", first.c_str(), part.c_str());
                     }
                 }
 

library/lua/dfhack.lua

@@ -386,15 +386,33 @@ end
 
 local internal = dfhack.internal
 
-internal.scripts = internal.scripts or {}
-internal.scriptPath = internal.scriptPath or {}
-internal.scriptMtime = internal.scriptMtime or {}
-internal.scriptRun = internal.scriptRun or {}
+Script = defclass(Script)
+function Script:init(path)
+    self.path = path
+    self.mtime = dfhack.filesystem.mtime(path)
+    self._flags = {}
+end
+function Script:get_flags()
+    if self.flags_mtime ~= dfhack.filesystem.mtime(self.path) then
+        self.flags_mtime = dfhack.filesystem.mtime(self.path)
+        self._flags = {}
+        local f = io.open(self.path)
+        local contents = f:read('*all')
+        f:close()
+        for line in contents:gmatch('%-%-@([^\n]+)') do
+            local chunk = load(line, self.path, 't', self._flags)
+            if chunk then
+                chunk()
+            else
+                dfhack.printerr('Parse error: ' .. line)
+            end
+        end
+    end
+    return self._flags
+end
 
+internal.scripts = internal.scripts or {}
 local scripts = internal.scripts
-local scriptPath = internal.scriptPath
-local scriptMtime = internal.scriptMtime
-local scriptRun = internal.scriptRun
 
 local hack_path = dfhack.getHackPath()
 
@@ -419,28 +437,69 @@ function dfhack.findScript(name)
     return nil
 end
 
+local valid_script_flags = {
+    enable = {required = true, error = 'Does not recognize enable/disable commands'},
+    enable_state = {required = false},
+    module = {
+        required = function(flags)
+            if flags.module_strict == false then return false end
+            return true
+        end,
+        error = 'Cannot be used as a module'
+    },
+    module_strict = {required = false},
+}
+
 function dfhack.run_script(name,...)
-    return dfhack.run_script_with_env(nil,name,...)
+    return dfhack.run_script_with_env(nil, name, nil, ...)
 end
 
+function dfhack.enable_script(name, state)
+    local res, err = dfhack.pcall(dfhack.run_script_with_env, nil, name, {enable=true, enable_state=state})
+    if not res then
+        dfhack.printerr(err.message)
+        qerror(('Cannot %s Lua script: %s'):format(state and 'enable' or 'disable', name))
+    end
+end
+
+function dfhack.reqscript(name)
+    _, env = dfhack.run_script_with_env(nil, name, {module=true})
+    return env
+end
+reqscript = dfhack.reqscript
+
 function dfhack.script_environment(name)
-    _, env = dfhack.run_script_with_env({moduleMode=true}, name)
+    _, env = dfhack.run_script_with_env(nil, name, {module=true, module_strict=false})
     return env
 end
 
-function dfhack.run_script_with_env(envVars,name,...)
+function dfhack.run_script_with_env(envVars, name, flags, ...)
+    if type(flags) ~= 'table' then flags = {} end
     local file = dfhack.findScript(name)
     if not file then
         error('Could not find script ' .. name)
     end
-    if scriptPath[name] and scriptPath[name] ~= file then
-        --new file path: must have loaded a different save or unloaded
-        scriptPath[name] = file
-        scriptMtime[file] = dfhack.filesystem.mtime(file)
-        --it is the responsibility of the script to clear its own data on unload so it's safe for us to not delete it here
+
+    if scripts[file] == nil then
+        scripts[file] = Script(file)
+    end
+    local script_flags = scripts[file]:get_flags()
+    for flag, value in pairs(flags) do
+        if value then
+            local v = valid_script_flags[flag]
+            if not v then
+                error('Invalid flag: ' .. flag)
+            elseif ((type(v.required) == 'boolean' and v.required) or
+                    (type(v.required) == 'function' and v.required(flags))) then
+                if not script_flags[flag] then
+                    local msg = v.error or 'Flag "' .. flag .. '" not recognized'
+                    error(name .. ': ' .. msg)
+                end
+            end
+        end
     end
 
-    local env = scripts[file]
+    local env = scripts[file].env
     if env == nil then
         env = {}
         setmetatable(env, { __index = base_env })
@@ -448,11 +507,13 @@ function dfhack.run_script_with_env(envVars,name,...)
     for x,y in pairs(envVars or {}) do
         env[x] = y
     end
+    env.dfhack_flags = flags
+    env.moduleMode = flags.module
     local f
     local perr
     local time = dfhack.filesystem.mtime(file)
-    if time == scriptMtime[file] and scriptRun[file] then
-        f = scriptRun[file]
+    if time == scripts[file].mtime and scripts[file].run then
+        f = scripts[file].run
     else
         --reload
         f, perr = loadfile(file, 't', env)
@@ -460,10 +521,10 @@ function dfhack.run_script_with_env(envVars,name,...)
             error(perr)
         end
         -- avoid updating mtime if the script failed to load
-        scriptMtime[file] = time
+        scripts[file].mtime = time
     end
-    scripts[file] = env
-    scriptRun[file] = f
+    scripts[file].env = env
+    scripts[file].run = f
     return f(...), env
 end
 

scripts/add-thought.lua

@@ -1,4 +1,5 @@
 -- Adds emotions to creatures.
+--@ module = true
 
 local utils=require('utils')
 

scripts/modtools/reaction-product-trigger.lua

@@ -1,6 +1,7 @@
 -- scripts/modtools/reaction-product-trigger.lua
 -- author expwnent
 -- trigger commands just before and after custom reactions produce items
+--@ module = true
 
 local eventful = require 'plugins.eventful'
 local utils = require 'utils'

scripts/modtools/reaction-trigger.lua

@@ -1,6 +1,7 @@
 -- scripts/modtools/reaction-trigger.lua
 -- author expwnent
 -- replaces autoSyndrome: trigger commands when custom reactions are completed
+--@ module = true
 
 local eventful = require 'plugins.eventful'
 local syndromeUtil = require 'syndrome-util'