Add DatabaseHandler class, clean up code, fix bugs

lib/commands.py

@@ -1,4 +1,3 @@
-import sqlite3
 import struct
 import re
 from lib import mapblock, blockfuncs, helpers
@@ -7,12 +6,11 @@ from lib import mapblock, blockfuncs, helpers
 # cloneblocks command
 #
 
-def clone_blocks(cursor, args):
+def clone_blocks(database, args):
     p1, p2 = helpers.args_to_mapblocks(args.p1, args.p2)
     offset = [n >> 4 for n in args.offset]
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, area=(p1, p2), inverse=args.inverse)
-            area=(p1, p2), inverse=args.inverse)
 
     # Sort the list to avoid overlapping of blocks when cloning.
     if offset[0] != 0: # Sort by x-value.
@@ -38,78 +36,66 @@ def clone_blocks(cursor, args):
             continue
         # Rehash the position.
         newPos = helpers.hash_pos(posVec)
-
+        # Get the mapblock and move it to the new location.
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        data = database.get_block(pos)
-        data = cursor.fetchone()[0]
+        database.set_block(newPos, data, force=True)
-        # Update mapblock or create a new one in new location.
-        cursor.execute(
-                "INSERT OR REPLACE INTO blocks (pos, data) VALUES (?, ?)",
-                (newPos, data))
 
 #
 # deleteblocks command
 #
 
-def delete_blocks(cursor, args):
+def delete_blocks(database, args):
     p1, p2 = helpers.args_to_mapblocks(args.p1, args.p2)
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(cursor, area=(p1, p2), inverse=args.inverse)
-            area=(p1, p2), inverse=args.inverse)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("DELETE FROM blocks WHERE pos = ?", (pos,))
+        database.delete_block(pos)
 
 #
 # fillblocks command
 #
 
-def fill_blocks(cursor, args):
+def fill_blocks(database, args):
     p1, p2 = helpers.args_to_mapblocks(args.p1, args.p2)
     name = bytes(args.replacename, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, area=(p1, p2), inverse=args.inverse)
-            area=(p1, p2), inverse=args.inverse)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         # Fill area with one type of node and delete everything else.
         parsedData.node_data = bytes(4096 * (parsedData.content_width +
                 parsedData.params_width))
-        parsedData.serialize_nimap([0], [name])
+        parsedData.serialize_nimap([name])
         parsedData.serialize_metadata([])
         parsedData.serialize_node_timers([])
 
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # overlayblocks command
 #
 
-def overlay_blocks(cursor, sCursor, args):
+def overlay_blocks(database, sDatabase, args):
     p1, p2 = helpers.args_to_mapblocks(args.p1, args.p2)
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(sCursor,
+    list = helpers.get_mapblocks(sDatabase,
             area=(p1, p2), inverse=args.inverse)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        sCursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        data = sDatabase.get_block(pos)
-        data = sCursor.fetchone()[0]
         # Update mapblock or create a new one in primary file.
-        cursor.execute(
+        database.set_block(pos, data, force=True)
-                "INSERT OR REPLACE INTO blocks (pos, data) VALUES (?, ?)",
-                (pos, data))
 
 #
 # replacenodes command
 #
 
-def replace_nodes(cursor, args):
+def replace_nodes(database, args):
     searchName = bytes(args.searchname, "utf-8")
     replaceName = bytes(args.replacename, "utf-8")
 
@@ -125,14 +111,11 @@ def replace_nodes(cursor, args):
         return
 
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName in nimap:
@@ -169,28 +152,23 @@ def replace_nodes(cursor, args):
                 nimap[nimap.index(searchName)] = replaceName
 
             parsedData.serialize_nimap(nimap)
-            data = parsedData.serialize()
+            database.set_block(pos, parsedData.serialize())
-
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # setparam2 command
 #
 
-def set_param2(cursor, args):
+def set_param2(database, args):
     if args.value < 0 or args.value > 255:
         helpers.throw_error("ERROR: param2 value must be between 0 and 255.")
 
     searchName = bytes(args.searchname, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName not in nimap:
@@ -200,30 +178,25 @@ def set_param2(cursor, args):
         bulkParam2 = bytearray(parsedData.node_data[12288:])
 
         for a in range(4096):
-            if parsedData.node_data[i * 2:(i + 1) * 2] == nodeId:
+            if parsedData.node_data[a * 2:(a + 1) * 2] == nodeId:
                 bulkParam2[a] = args.value
 
         parsedData.node_data = (parsedData.node_data[:12288] +
                 bytes(bulkParam2))
-
+        database.set_block(pos, parsedData.serialize())
-        data = parsedData.serialize()
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # deletemeta command
 #
 
-def delete_meta(cursor, args):
+def delete_meta(database, args):
     searchName = bytes(args.searchname, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName not in nimap:
@@ -238,26 +211,22 @@ def delete_meta(cursor, args):
                 del metaList[a]
 
         parsedData.serialize_metadata(metaList)
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # setmetavar command
 #
 
-def set_meta_var(cursor, args):
+def set_meta_var(database, args):
     searchName = bytes(args.searchname, "utf-8")
     key = bytes(args.key, "utf-8")
     value = bytes(args.value, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName not in nimap:
@@ -279,26 +248,22 @@ def set_meta_var(cursor, args):
                         parsedData.metadata_version)
 
         parsedData.serialize_metadata(metaList)
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # replaceininv command
 #
 
-def replace_in_inv(cursor, args):
+def replace_in_inv(database, args):
     searchName = bytes(args.searchname, "utf-8")
     searchItem = bytes(args.searchitem, "utf-8")
     replaceItem = bytes(args.replaceitem, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName not in nimap:
@@ -328,27 +293,24 @@ def replace_in_inv(cursor, args):
 
                     invList[b] = b" ".join(splitItem)
 
+                # Re-join node inventory.
                 metaList[a]["inv"] = b"\n".join(invList)
 
         parsedData.serialize_metadata(metaList)
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # deletetimers
 #
 
-def delete_timers(cursor, args):
+def delete_timers(database, args):
     searchName = bytes(args.searchname, "utf-8")
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
         nimap = parsedData.deserialize_nimap()
 
         if searchName not in nimap:
@@ -364,30 +326,26 @@ def delete_timers(cursor, args):
                 del timers[a]
 
         parsedData.serialize_node_timers(timers)
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))
 
 #
 # deleteobjects
 #
 
-def delete_objects(cursor, args):
+def delete_objects(database, args):
     if args.item:
         searchName = b"__builtin:item"
     else:
         searchName = bytes(args.searchname, "utf-8")
 
     progress = helpers.Progress()
-    list = helpers.get_mapblocks(cursor,
+    list = helpers.get_mapblocks(database, name=searchName)
-            name=searchName)
     itemstringFormat = re.compile(
             b"\[\"itemstring\"\] = \"(?P<name>[a-zA-Z0-9_:]+)")
 
     for i, pos in enumerate(list):
         progress.print_bar(i, len(list))
-        cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        parsedData = mapblock.MapBlock(database.get_block(pos))
-        data = cursor.fetchone()[0]
-        parsedData = mapblock.MapBlock(data)
 
         if parsedData.static_object_count == 0:
             continue
@@ -408,5 +366,4 @@ def delete_objects(cursor, args):
                     del objects[a]
 
         parsedData.serialize_static_objects(objects)
-        data = parsedData.serialize()
+        database.set_block(pos, parsedData.serialize())
-        cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?", (data, pos))

lib/helpers.py

@@ -1,6 +1,55 @@
-import sys
+import sqlite3
 import math
 import time
+import sys
+
+class DatabaseHandler:
+    """Handles the Sqlite database and provides useful functions."""
+
+    def __init__(self, filename, type):
+        if not filename:
+            throw_error("Please specify a map file ({:s}).".format(type))
+
+        try:
+            tempFile = open(filename, 'r')
+            tempFile.close()
+        except:
+            throw_error("Map file does not exist ({:s}).".format(type))
+
+        self.database = sqlite3.connect(filename)
+        self.cursor = self.database.cursor()
+
+        try:
+            self.cursor.execute("SELECT * FROM blocks")
+        except sqlite3.DatabaseError:
+            throw_error("File is not a valid map file ({:s}).".format(type))
+
+
+    def get_block(self, pos):
+        self.cursor.execute("SELECT data FROM blocks WHERE pos = ?", (pos,))
+        return self.cursor.fetchone()[0]
+
+
+    def delete_block(self, pos):
+        self.cursor.execute("DELETE FROM blocks WHERE pos = ?", (pos,))
+
+
+    def set_block(self, pos, data, force = False):
+        if force:
+            self.cursor.execute(
+                    "INSERT OR REPLACE INTO blocks (pos, data) VALUES (?, ?)",
+                    (pos, data))
+        else:
+            self.cursor.execute("UPDATE blocks SET data = ? WHERE pos = ?",
+                    (data, pos))
+
+
+    def close(self, commit = False):
+        if commit:
+            self.database.commit()
+
+        self.database.close()
+
 
 class Progress:
     """Prints a progress bar with time elapsed."""
@@ -101,12 +150,12 @@ def is_in_range(num, area):
     return True
 
 
-def get_mapblocks(cursor, area = None, name = None, inverse = False):
+def get_mapblocks(database, area = None, name = None, inverse = False):
     batch = []
     list = []
 
     while True:
-        batch = cursor.fetchmany(1000)
+        batch = database.cursor.fetchmany(1000)
         # Exit if we run out of database entries.
         if len(batch) == 0:
             break

mapedit.py

@@ -1,6 +1,5 @@
-import sys
 import argparse
-import sqlite3
+import sys
 import re
 from lib import commands, helpers
 
@@ -79,14 +78,14 @@ parser_deletemeta = subparsers.add_parser("deletemeta",
         help="Delete metadata of all of a certain type of node.")
 parser_deletemeta.set_defaults(func=commands.delete_meta)
 
-parser_replaceininv = subparsers.add_parser("replaceininv",
-        help="Replace one item with another in inventories certain nodes.")
-parser_replaceininv.set_defaults(func=commands.replace_in_inv)
-
 parser_setmetavar = subparsers.add_parser("setmetavar",
         help="Set a value in the metadata of all of a certain type of node.")
 parser_setmetavar.set_defaults(func=commands.set_meta_var)
 
+parser_replaceininv = subparsers.add_parser("replaceininv",
+        help="Replace one item with another in inventories certain nodes.")
+parser_replaceininv.set_defaults(func=commands.replace_in_inv)
+
 parser_deletetimers = subparsers.add_parser("deletetimers",
         help="Delete node timers of all of a certain type of node.")
 parser_deletetimers.set_defaults(func=commands.delete_timers)
@@ -153,18 +152,10 @@ for param in ("searchname", "replacename", "searchitem", "replaceitem"):
 
         if (nameFormat.match(value) == None and value != "air" and not
                 (param == "replaceitem" and value == "Empty")):
-            helpers.throw_error("Invalid node name.")
+            helpers.throw_error("Invalid node name ({:s}).".format(param))
 
-helpers.verify_file(args.f, "Primary map file does not exist.")
+# Attempt to open database.
-
+db = helpers.DatabaseHandler(args.f, "primary")
-db = sqlite3.connect(args.f)
-cursor = db.cursor()
-
-# Test for database validity.
-try:
-    cursor.execute("SELECT * FROM blocks")
-except sqlite3.DatabaseError:
-    helpers.throw_error("Primary map file is not a valid map database.")
 
 if not args.silencewarnings and input(
         "WARNING: Using this tool can potentially cause permanent\n"
@@ -175,31 +166,17 @@ if not args.silencewarnings and input(
     sys.exit()
 
 if args.command == "overlayblocks":
-    if not args.s:
-        helpers.throw_error("Command requires a secondary map file.")
-
     if args.s == args.f:
         helpers.throw_error("Primary and secondary map files are the same.")
 
-    helpers.verify_file(args.s, "Secondary map file does not exist.")
+    sDb = helpers.DatabaseHandler(args.s, "secondary")
 
-    sDb = sqlite3.connect(args.s)
+    args.func(db, sDb, args)
-    sCursor = sDb.cursor()
-
-    # Test for database validity.
-    try:
-        sCursor.execute("SELECT * FROM blocks")
-    except sqlite3.DatabaseError:
-        helpers.throw_error("Secondary map file is not a valid map database.")
-
-    args.func(cursor, sCursor, args)
     sDb.close()
 else:
-    args.func(cursor, args)
+    args.func(db, args)
 
 print("\nSaving file...")
-
+db.close(commit=True)
-db.commit()
-db.close()
 
 print("Done.")