[drain-aquifer] more control over which levels get drained

changelog.txt

@@ -29,6 +29,9 @@ Template for new versions:
 ## New Tools
 
 ## New Features
+- `drain-aquifer`: gained ability to drain just above or below a certain z-level
+- `drain-aquifer`: new option to drain all layers except for the first N aquifer layers, in case you want some aquifer layers but not too many
+- `gui/control-panel`: ``drain-aquifer --top 2`` added as an autostart option
 
 ## Fixes
 - `suspendmanager`: fix errors when constructing near the map edge

docs/drain-aquifer.rst

@@ -2,14 +2,40 @@ drain-aquifer
 =============
 
 .. dfhack-tool::
-    :summary: Remove all aquifers on the map.
+    :summary: Remove some or all aquifers on the map.
     :tags: fort armok map
 
-This tool irreversibly removes all 'aquifer' tags from the map blocks.
+This tool irreversibly removes 'aquifer' tags from map blocks. Also see
+`prospect` for discovering the range of layers that currently have aquifers.
 
 Usage
 -----
 
 ::
 
-    drain-aquifer
+    drain-aquifer [<options>]
+
+Examples
+--------
+
+``drain-aquifer``
+    Remove all aquifers on the map.
+``drain-aquifer --top 2``
+    Remove all aquifers on the map except for the top 2 levels of aquifer.
+``drain-aquifer -d``
+    Remove all aquifers on the current z-level and below.
+
+Options
+-------
+
+``-t``, ``--top <num>``
+    Remove all aquifers on the map except for the top ``<num>`` levels,
+    starting from the first level that has an aquifer tile. Note that there may
+    be less than ``<num>`` levels of aquifer after the command is run if the
+    levels of aquifer are not contiguous.
+``-d``, ``--zdown``
+    Remove all aquifers on the current z-level and below.
+``-u``, ``--zup``
+    Remove all aquifers on the current z-level and above.
+``-z``, ``--cur-zlevel``
+    Remove all aquifers on the current z-level.

drain-aquifer.lua

@@ -1,40 +1,68 @@
--- Remove all aquifers from the map
---[====[
+local argparse = require('argparse')
 
-drain-aquifer
-=============
-Remove all 'aquifer' tags from the map blocks. Irreversible.
-
-]====]
+local zmin = 0
+local zmax = df.global.world.map.z_count - 1
 
 local function drain()
     local layers = {} --as:bool[]
     local layer_count = 0
     local tile_count = 0
 
-    for k, block in ipairs(df.global.world.map.map_blocks) do
-        if block.flags.has_aquifer then
-            block.flags.has_aquifer = false
-            block.flags.check_aquifer = false
-
-            for x, row in ipairs(block.designation) do
-                for y, tile in ipairs(row) do
-                    if tile.water_table then
-                        tile.water_table = false
-                        tile_count = tile_count + 1
-                    end
+    for _, block in ipairs(df.global.world.map.map_blocks) do
+        if not block.flags.has_aquifer then goto continue end
+        if block.map_pos.z < zmin or block.map_pos.z > zmax then goto continue end
+
+        block.flags.has_aquifer = false
+        block.flags.check_aquifer = false
+
+        for _, row in ipairs(block.designation) do
+            for _, tile in ipairs(row) do
+                if tile.water_table then
+                    tile.water_table = false
+                    tile_count = tile_count + 1
                 end
             end
+        end
 
-            if not layers[block.map_pos.z] then
-                layers[block.map_pos.z] = true
-                layer_count = layer_count + 1
-            end
+        if not layers[block.map_pos.z] then
+            layers[block.map_pos.z] = true
+            layer_count = layer_count + 1
         end
+        ::continue::
     end
 
-    print("Cleared "..tile_count.." aquifer tile"..((tile_count ~= 1) and "s" or "")..
-        " in "..layer_count.." layer"..((layer_count ~= 1) and "s" or "")..".")
+    print(('Cleared %d aquifer tile%s in %d layer%s.'):format(
+        tile_count, (tile_count ~= 1) and 's' or '', layer_count, (layer_count ~= 1) and 's' or ''))
+end
+
+local help = false
+local top = 0
+
+local positionals = argparse.processArgsGetopt({...}, {
+    {'h', 'help', handler=function() help = true end},
+    {'t', 'top', hasArg=true, handler=function(optarg) top = argparse.nonnegativeInt(optarg, 'top') end},
+    {'d', 'zdown', handler=function() zmax = df.global.window_z end},
+    {'u', 'zup', handler=function() zmin = df.global.window_z end},
+    {'z', 'cur-zlevel', handler=function() zmax, zmin = df.global.window_z, df.global.window_z end},
+})
+
+if help or positionals[1] == 'help' then
+    print(dfhack.script_help())
+    return
+end
+
+if top > 0 then
+    zmax = -1
+    for _, block in ipairs(df.global.world.map.map_blocks) do
+        if block.flags.has_aquifer and zmax < block.map_pos.z then
+            zmax = block.map_pos.z
+        end
+    end
+    zmax = zmax - top
+    if zmax < zmin then
+        print('No aquifer levels need draining')
+        return
+    end
 end
 
-drain(...)
+drain()

gui/control-panel.lua

@@ -46,6 +46,7 @@ local FORT_AUTOSTART = {
     'ban-cooking all',
     'buildingplan set boulders false',
     'buildingplan set logs false',
+    'drain-aquifer --top 2',
     'fix/blood-del fort',
     'light-aquifers-only fort',
 }