Merge remote-tracking branch 'myk002/myk_blueprint2' into develop

plugins/blueprint.cpp

@@ -10,6 +10,7 @@
 
 #include "Console.h"
 #include "DataDefs.h"
+#include "DataFuncs.h"
 #include "DataIdentity.h"
 #include "LuaTools.h"
 #include "PluginManager.h"
@@ -85,27 +86,27 @@ static const struct_field_info blueprint_options_fields[] = {
 };
 struct_identity blueprint_options::_identity(sizeof(blueprint_options), &df::allocator_fn<blueprint_options>, NULL, "blueprint_options", NULL, blueprint_options_fields);
 
-command_result blueprint(color_ostream &out, vector<string> &parameters);
+command_result blueprint(color_ostream &, vector<string> &);
 
-DFhackCExport command_result plugin_init(color_ostream &out, vector<PluginCommand> &commands)
+DFhackCExport command_result plugin_init(color_ostream &, vector<PluginCommand> &commands)
 {
     commands.push_back(PluginCommand("blueprint", "Record the structure of a live game map in a quickfort blueprint", blueprint, false));
     return CR_OK;
 }
 
-DFhackCExport command_result plugin_shutdown(color_ostream &out)
+DFhackCExport command_result plugin_shutdown(color_ostream &)
 {
     return CR_OK;
 }
 
-pair<uint32_t, uint32_t> get_building_size(df::building* b)
+static pair<uint32_t, uint32_t> get_building_size(df::building* b)
 {
     return pair<uint32_t, uint32_t>(b->x2 - b->x1 + 1, b->y2 - b->y1 + 1);
 }
 
-char get_tile_dig(int32_t x, int32_t y, int32_t z)
+static char get_tile_dig(int32_t x, int32_t y, int32_t z)
 {
-    df::tiletype *tt = Maps::getTileType(x, y , z);
+    df::tiletype *tt = Maps::getTileType(x, y, z);
     df::tiletype_shape ts = tileShape(tt ? *tt : tiletype::Void);
     switch (ts)
     {
@@ -132,7 +133,7 @@ char get_tile_dig(int32_t x, int32_t y, int32_t z)
     }
 }
 
-string get_tile_build(uint32_t x, uint32_t y, df::building* b)
+static string get_tile_build(uint32_t x, uint32_t y, df::building* b)
 {
     if (! b)
         return " ";
@@ -513,7 +514,7 @@ string get_tile_build(uint32_t x, uint32_t y, df::building* b)
     }
 }
 
-string get_tile_place(uint32_t x, uint32_t y, df::building* b)
+static string get_tile_place(uint32_t x, uint32_t y, df::building* b)
 {
     if (! b || b->getType() != building_type::Stockpile)
         return " ";
@@ -579,24 +580,28 @@ string get_tile_place(uint32_t x, uint32_t y, df::building* b)
     return out.str();
 }
 
-string get_tile_query(df::building* b)
+static string get_tile_query(df::building* b)
 {
     if (b && b->is_room)
         return "r+";
     return " ";
 }
 
-void init_stream(ofstream &out, string basename, string target)
+// returns filename
+static string init_stream(ofstream &out, string basename, string target)
 {
     std::ostringstream out_path;
     out_path << basename << "-" << target << ".csv";
-    out.open(out_path.str(), ofstream::trunc);
+    string path = out_path.str();
+    out.open(path, ofstream::trunc);
     out << "#" << target << endl;
+    return path;
 }
 
-command_result do_transform(const DFCoord &start, const DFCoord &end,
+static bool do_transform(const DFCoord &start, const DFCoord &end,
-                            const blueprint_options &options,
+                         const blueprint_options &options,
-                            std::ostringstream &err)
+                         vector<string> &files,
+                         std::ostringstream &err)
 {
     ofstream dig, build, place, query;
 
@@ -609,24 +614,24 @@ command_result do_transform(const DFCoord &start, const DFCoord &end,
     if (!Filesystem::mkdir_recursive(parent_path))
     {
         err << "could not create output directory: '" << parent_path << "'";
-        return CR_FAILURE;
+        return false;
     }
 
-    if (options.auto_phase || options.query)
+    if (options.auto_phase || options.dig)
     {
-        init_stream(query, basename, "query");
+        files.push_back(init_stream(dig, basename, "dig"));
     }
-    if (options.auto_phase || options.place)
+    if (options.auto_phase || options.build)
     {
-        init_stream(place, basename, "place");
+        files.push_back(init_stream(build, basename, "build"));
     }
-    if (options.auto_phase || options.build)
+    if (options.auto_phase || options.place)
     {
-        init_stream(build, basename, "build");
+        files.push_back(init_stream(place, basename, "place"));
     }
-    if (options.auto_phase || options.dig)
+    if (options.auto_phase || options.query)
     {
-        init_stream(dig, basename, "dig");
+        files.push_back(init_stream(query, basename, "query"));
     }
 
     const int32_t z_inc = start.z < end.z ? 1 : -1;
@@ -677,14 +682,14 @@ command_result do_transform(const DFCoord &start, const DFCoord &end,
     if (options.auto_phase || options.dig)
         dig.close();
 
-    return CR_OK;
+    return true;
 }
 
-static bool get_options(blueprint_options &opts,
+static bool get_options(color_ostream &out,
+                        blueprint_options &opts,
                         const vector<string> &parameters)
 {
     auto L = Lua::Core::State;
-    color_ostream_proxy out(Core::getInstance().getConsole());
     Lua::StackUnwinder top(L);
 
     if (!lua_checkstack(L, parameters.size() + 2) ||
@@ -706,10 +711,9 @@ static bool get_options(blueprint_options &opts,
     return true;
 }
 
-static void print_help()
+static void print_help(color_ostream &out)
 {
     auto L = Lua::Core::State;
-    color_ostream_proxy out(Core::getInstance().getConsole());
     Lua::StackUnwinder top(L);
 
     if (!lua_checkstack(L, 1) ||
@@ -720,7 +724,11 @@ static void print_help()
     }
 }
 
-command_result blueprint(color_ostream &out, vector<string> &parameters)
+// returns whether blueprint generation was successful. populates files with the
+// names of the files that were generated
+static bool do_blueprint(color_ostream &out,
+                         const vector<string> &parameters,
+                         vector<string> &files)
 {
     CoreSuspender suspend;
 
@@ -740,16 +748,16 @@ command_result blueprint(color_ostream &out, vector<string> &parameters)
     }
 
     blueprint_options options;
-    if (!get_options(options, parameters) || options.help)
+    if (!get_options(out, options, parameters) || options.help)
     {
-        print_help();
+        print_help(out);
-        return options.help ? CR_OK : CR_FAILURE;
+        return options.help;
     }
 
     if (!Maps::IsValid())
     {
         out.printerr("Map is not available!\n");
-        return CR_FAILURE;
+        return false;
     }
 
     // start coordinates can come from either the commandline or the map cursor
@@ -760,14 +768,14 @@ command_result blueprint(color_ostream &out, vector<string> &parameters)
         {
             out.printerr("Can't get cursor coords! Make sure you specify the"
                     " --cursor parameter or have an active cursor in DF.\n");
-            return CR_FAILURE;
+            return false;
         }
     }
     if (!Maps::isValidTilePos(start))
     {
         out.printerr("Invalid start position: %d,%d,%d\n",
                      start.x, start.y, start.z);
-        return CR_FAILURE;
+        return false;
     }
 
     // end coords are one beyond the last processed coordinate. note that
@@ -789,8 +797,55 @@ command_result blueprint(color_ostream &out, vector<string> &parameters)
         end.z = -1;
 
     std::ostringstream err;
-    command_result result = do_transform(start, end, options, err);
+    if (!do_transform(start, end, options, files, err))
-    if (result != CR_OK)
+    {
         out.printerr("%s\n", err.str().c_str());
-    return result;
+        return false;
+    }
+    return true;
 }
+
+// entrypoint when called from Lua. returns the names of the generated files
+static int run(lua_State *L)
+{
+    int argc = lua_gettop(L);
+    vector<string> argv;
+
+    for (int i = 1; i <= argc; ++i)
+    {
+        const char *s = lua_tostring(L, i);
+        if (s == NULL)
+            luaL_error(L, "all parameters must be strings");
+        argv.push_back(s);
+    }
+
+    vector<string> files;
+    color_ostream *out = Lua::GetOutput(L);
+    if (!out)
+        out = &Core::getInstance().getConsole();
+    if (do_blueprint(*out, argv, files))
+    {
+        Lua::PushVector(L, files);
+        return 1;
+    }
+
+    return 0;
+}
+
+command_result blueprint(color_ostream &out, vector<string> &parameters)
+{
+    vector<string> files;
+    if (do_blueprint(out, parameters, files))
+    {
+        out.print("Generated blueprint file(s):\n");
+        for (string &fname : files)
+            out.print("  %s\n", fname.c_str());
+        return CR_OK;
+    }
+    return CR_FAILURE;
+}
+
+DFHACK_PLUGIN_LUA_COMMANDS {
+    DFHACK_LUA_COMMAND(run),
+    DFHACK_LUA_END
+};

plugins/lua/blueprint.lua

@@ -138,10 +138,8 @@ function parse_commandline(opts, ...)
     parse_positionals(opts, positionals, depth and 4 or 3)
 end
 
--- compatibility with old exported API. we route the request back through
+-- compatibility with old exported API.
--- run_command so we have a unified path for parameter processing and invariant
+local function do_phase(start_pos, end_pos, name, phase)
--- checking.
-local function do_blueprint(start_pos, end_pos, name, phase)
     local width = math.abs(start_pos.x - end_pos.x) + 1
     local height = math.abs(start_pos.y - end_pos.y) + 1
     local depth = math.abs(start_pos.z - end_pos.z) + 1
@@ -153,13 +151,11 @@ local function do_blueprint(start_pos, end_pos, name, phase)
 
     local cursor = ('--cursor=%d,%d,%d'):format(x, y, z)
 
-    return dfhack.run_command('blueprint',
+    run(tostring(width), tostring(height), tostring(depth), tostring(name),
-                              tostring(width), tostring(height),
+        phase, cursor)
-                              tostring(depth), tostring(name),
-                              phase, cursor)
 end
 for phase in pairs(valid_phases) do
-    _ENV[phase] = function(s, e, n) do_blueprint(s, e, n, phase) end
+    _ENV[phase] = function(s, e, n) do_phase(s, e, n, phase) end
 end
 
 return _ENV

test/plugins/blueprint.lua

@@ -121,44 +121,44 @@ function test.parse_commandline()
                        'zero depth')
 end
 
-function test.do_blueprint_positive_dims()
+function test.do_phase_positive_dims()
-    local mock_run_command = mock.func()
+    local mock_run = mock.func()
-    mock.patch(dfhack, 'run_command', mock_run_command,
+    mock.patch(b, 'run', mock_run,
         function()
               local spos = {x=10, y=20, z=30}
               local epos = {x=11, y=21, z=31}
               b.query(spos, epos, 'imaname')
-              expect.eq(1, mock_run_command.call_count)
+              expect.eq(1, mock_run.call_count)
-              expect.table_eq({'blueprint', '2', '2', '2', 'imaname', 'query',
+              expect.table_eq({'2', '2', '2', 'imaname', 'query',
                                '--cursor=10,20,30'},
-                              mock_run_command.call_args[1])
+                              mock_run.call_args[1])
         end)
 end
 
-function test.do_blueprint_negative_dims()
+function test.do_phase_negative_dims()
-    local mock_run_command = mock.func()
+    local mock_run = mock.func()
-    mock.patch(dfhack, 'run_command', mock_run_command,
+    mock.patch(b, 'run', mock_run,
         function()
               local spos = {x=11, y=21, z=31}
               local epos = {x=10, y=20, z=30}
               b.query(spos, epos, 'imaname')
-              expect.eq(1, mock_run_command.call_count)
+              expect.eq(1, mock_run.call_count)
-              expect.table_eq({'blueprint', '2', '2', '-2', 'imaname', 'query',
+              expect.table_eq({'2', '2', '-2', 'imaname', 'query',
                                '--cursor=10,20,31'},
-                              mock_run_command.call_args[1])
+                              mock_run.call_args[1])
         end)
 end
 
-function test.do_blueprint_ensure_cursor_is_at_upper_left()
+function test.do_phase_ensure_cursor_is_at_upper_left()
-    local mock_run_command = mock.func()
+    local mock_run = mock.func()
-    mock.patch(dfhack, 'run_command', mock_run_command,
+    mock.patch(b, 'run', mock_run,
         function()
               local spos = {x=11, y=20, z=30}
               local epos = {x=10, y=21, z=31}
               b.query(spos, epos, 'imaname')
-              expect.eq(1, mock_run_command.call_count)
+              expect.eq(1, mock_run.call_count)
-              expect.table_eq({'blueprint', '2', '2', '2', 'imaname', 'query',
+              expect.table_eq({'2', '2', '2', 'imaname', 'query',
                                '--cursor=10,20,30'},
-                              mock_run_command.call_args[1])
+                              mock_run.call_args[1])
         end)
 end