Refactor code functions

Refactor the code functions to interact directly with the NBT classes instead of a weird format

Author: gentlegiantJGC

PyMCTranslate/build_number.json

@@ -1 +1 @@
-367
+368

PyMCTranslate/code_functions/banner_pattern_2u.py

@@ -1,21 +1,26 @@
+from amulet_nbt import CompoundTag, ListTag, IntTag
+
 def main(nbt):
-    if (
-        nbt[0] == "compound"
-        and "Patterns" in nbt[1]
-        and nbt[1]["Patterns"][0] == "list"
-    ):
-        return [
-            [
-                "",
-                "compound",
-                [("utags", "compound"), ("Patterns", "list"), (index, "compound")],
-                "Color",
-                ["int", 15 - pattern[1]["Color"][1]],
-            ]
-            for index, pattern in enumerate(nbt[1]["Patterns"][1])
-            if pattern[0] == "compound"
-            and "Color" in pattern[1]
-            and pattern[1]["Color"][0] == "int"
-        ]
-    else:
+    if not isinstance(nbt, CompoundTag):
         return []
+    patterns = nbt.get("Patterns")
+    if not isinstance(patterns, ListTag):
+        return []
+
+    tags = []
+    index = 0
+    for pattern in patterns:
+        if not isinstance(pattern, CompoundTag):
+            continue
+        colour = pattern.get("Color")
+        if not isinstance(colour, IntTag):
+            continue
+        tags.append([
+            "",
+            "compound",
+            [("utags", "compound"), ("Patterns", "list"), (index, "compound")],
+            "Color",
+            IntTag(15 - colour.py_int),
+        ])
+        index += 1
+    return tags

PyMCTranslate/code_functions/banner_pattern_fu.py

@@ -1,23 +1,29 @@
+from amulet_nbt import CompoundTag, ListTag, IntTag
+
+
 def main(nbt):
-    if (
-        nbt[0] == "compound"
-        and "utags" in nbt[1]
-        and nbt[1]["utags"][0] == "compound"
-        and "Patterns" in nbt[1]["utags"][1]
-        and nbt[1]["utags"][1]["Patterns"][0] == "list"
-    ):
-        return [
-            [
-                "",
-                "compound",
-                [("Patterns", "list"), (index, "compound")],
-                "Color",
-                ["int", 15 - pattern[1]["Color"][1]],
-            ]
-            for index, pattern in enumerate(nbt[1]["utags"][1]["Patterns"][1])
-            if pattern[0] == "compound"
-            and "Color" in pattern[1]
-            and pattern[1]["Color"][0] == "int"
-        ]
-    else:
+    if not isinstance(nbt, CompoundTag):
         return []
+    utags = nbt.get("utags")
+    if not isinstance(utags, CompoundTag):
+        return []
+    patterns = utags.get("Patterns")
+    if not isinstance(patterns, ListTag):
+        return []
+    tags = []
+    index = 0
+    for pattern in patterns:
+        if not isinstance(pattern, CompoundTag):
+            continue
+        colour = pattern.get("Color")
+        if not isinstance(colour, IntTag):
+            continue
+        tags.append([
+            "",
+            "compound",
+            [("Patterns", "list"), (index, "compound")],
+            "Color",
+            IntTag(15 - colour.py_int),
+        ])
+        index += 1
+    return tags

PyMCTranslate/code_functions/bedrock_chest_connection_other_left.py

@@ -1,31 +1,35 @@
-def main(nbt, properties, location):
-    if (
-        nbt[0] == "compound"
-        and "pairlead" in nbt[1]
-        and nbt[1]["pairlead"][0] == "byte"
-        and nbt[1]["pairlead"][1] == 1
-        and "pairx" in nbt[1]
-        and nbt[1]["pairx"][0] == "int"
-        and "pairz" in nbt[1]
-        and nbt[1]["pairz"][0] == "int"
-        and "facing_direction" in properties
-    ):
-        dx = nbt[1]["pairx"][1] - location[0]
-        dz = nbt[1]["pairz"][1] - location[2]
-        if properties["facing_direction"].py_int == 2:  # north
-            if dz == 0:
-                if dx == -1:
-                    return {"connection": "left"}
-        elif properties["facing_direction"].py_int == 3:  # south
-            if dz == 0:
-                if dx == 1:
-                    return {"connection": "left"}
-        elif properties["facing_direction"].py_int == 4:  # west
-            if dx == 0:
-                if dz == 1:
-                    return {"connection": "left"}
-        elif properties["facing_direction"].py_int == 5:  # east
-            if dx == 0:
-                if dz == -1:
-                    return {"connection": "left"}
+from typing import Union
+
+from amulet_nbt import CompoundTag, ByteTag, ShortTag, IntTag, LongTag, StringTag
+
+
+def main(nbt, properties: dict[str, Union[ByteTag, ShortTag, IntTag, LongTag, StringTag]], location: tuple[int, int, int]) -> dict[str, str]:
+    if not isinstance(nbt, CompoundTag):
+        return {}
+    facing_direction = properties.get("facing_direction")
+    if not isinstance(facing_direction, IntTag):
+        return {}
+    pairlead = nbt.get("pairlead")
+    if isinstance(pairlead, ByteTag) and pairlead.py_int == 1:
+        pair_x = nbt.get("pairx")
+        if not isinstance(pair_x, IntTag):
+            return {}
+        pair_z = nbt.get("pairz")
+        if not isinstance(pair_z, IntTag):
+            return {}
+
+        dx = pair_x.py_int - location[0]
+        dz = pair_z.py_int - location[2]
+        if facing_direction.py_int == 2:  # north
+            if dz == 0 and dx == -1:
+                return {"connection": "left"}
+        elif facing_direction.py_int == 3:  # south
+            if dz == 0 and dx == 1:
+                return {"connection": "left"}
+        elif facing_direction.py_int == 4:  # west
+            if dx == 0 and dz == 1:
+                return {"connection": "left"}
+        elif facing_direction.py_int == 5:  # east
+            if dx == 0 and dz == -1:
+                return {"connection": "left"}
     return {}

PyMCTranslate/code_functions/bedrock_chest_connection_other_left_120.py

@@ -1,31 +1,35 @@
-def main(nbt, properties, location):
-    if (
-        nbt[0] == "compound"
-        and "pairlead" in nbt[1]
-        and nbt[1]["pairlead"][0] == "byte"
-        and nbt[1]["pairlead"][1] == 1
-        and "pairx" in nbt[1]
-        and nbt[1]["pairx"][0] == "int"
-        and "pairz" in nbt[1]
-        and nbt[1]["pairz"][0] == "int"
-        and "minecraft:cardinal_direction" in properties
-    ):
-        dx = nbt[1]["pairx"][1] - location[0]
-        dz = nbt[1]["pairz"][1] - location[2]
-        if properties["minecraft:cardinal_direction"].py_str == "north":  # north
-            if dz == 0:
-                if dx == -1:
-                    return {"connection": "left"}
-        elif properties["minecraft:cardinal_direction"].py_str == "south":  # south
-            if dz == 0:
-                if dx == 1:
-                    return {"connection": "left"}
-        elif properties["minecraft:cardinal_direction"].py_str == "west":  # west
-            if dx == 0:
-                if dz == 1:
-                    return {"connection": "left"}
-        elif properties["minecraft:cardinal_direction"].py_str == "east":  # east
-            if dx == 0:
-                if dz == -1:
-                    return {"connection": "left"}
+from typing import Union
+
+from amulet_nbt import CompoundTag, ByteTag, ShortTag, IntTag, LongTag, StringTag
+
+
+def main(nbt, properties: dict[str, Union[ByteTag, ShortTag, IntTag, LongTag, StringTag]], location: tuple[int, int, int]) -> dict[str, str]:
+    if not isinstance(nbt, CompoundTag):
+        return {}
+    facing_direction = properties.get("minecraft:cardinal_direction")
+    if not isinstance(facing_direction, StringTag):
+        return {}
+    pairlead = nbt.get("pairlead")
+    if isinstance(pairlead, ByteTag) and pairlead.py_int == 1:
+        pair_x = nbt.get("pairx")
+        if not isinstance(pair_x, IntTag):
+            return {}
+        pair_z = nbt.get("pairz")
+        if not isinstance(pair_z, IntTag):
+            return {}
+
+        dx = pair_x.py_int - location[0]
+        dz = pair_z.py_int - location[2]
+        if facing_direction.py_str == "north":  # north
+            if dz == 0 and dx == -1:
+                return {"connection": "left"}
+        elif facing_direction.py_str == "south":  # south
+            if dz == 0 and dx == 1:
+                return {"connection": "left"}
+        elif facing_direction.py_str == "west":  # west
+            if dx == 0 and dz == 1:
+                return {"connection": "left"}
+        elif facing_direction.py_str == "east":  # east
+            if dx == 0 and dz == -1:
+                return {"connection": "left"}
     return {}

PyMCTranslate/code_functions/bedrock_chest_connection_other_right.py

@@ -1,31 +1,35 @@
-def main(nbt, properties, location):
-    if (
-        nbt[0] == "compound"
-        and "pairlead" in nbt[1]
-        and nbt[1]["pairlead"][0] == "byte"
-        and nbt[1]["pairlead"][1] == 1
-        and "pairx" in nbt[1]
-        and nbt[1]["pairx"][0] == "int"
-        and "pairz" in nbt[1]
-        and nbt[1]["pairz"][0] == "int"
-        and "facing_direction" in properties
-    ):
-        dx = nbt[1]["pairx"][1] - location[0]
-        dz = nbt[1]["pairz"][1] - location[2]
-        if properties["facing_direction"].py_int == 2:  # north
-            if dz == 0:
-                if dx == 1:
-                    return {"connection": "right"}
-        elif properties["facing_direction"].py_int == 3:  # south
-            if dz == 0:
-                if dx == -1:
-                    return {"connection": "right"}
-        elif properties["facing_direction"].py_int == 4:  # west
-            if dx == 0:
-                if dz == -1:
-                    return {"connection": "right"}
-        elif properties["facing_direction"].py_int == 5:  # east
-            if dx == 0:
-                if dz == 1:
-                    return {"connection": "right"}
+from typing import Union
+
+from amulet_nbt import CompoundTag, ByteTag, ShortTag, IntTag, LongTag, StringTag
+
+
+def main(nbt, properties: dict[str, Union[ByteTag, ShortTag, IntTag, LongTag, StringTag]], location: tuple[int, int, int]) -> dict[str, str]:
+    if not isinstance(nbt, CompoundTag):
+        return {}
+    facing_direction = properties.get("facing_direction")
+    if not isinstance(facing_direction, IntTag):
+        return {}
+    pairlead = nbt.get("pairlead")
+    if isinstance(pairlead, ByteTag) and pairlead.py_int == 1:
+        pair_x = nbt.get("pairx")
+        if not isinstance(pair_x, IntTag):
+            return {}
+        pair_z = nbt.get("pairz")
+        if not isinstance(pair_z, IntTag):
+            return {}
+
+        dx = pair_x.py_int - location[0]
+        dz = pair_z.py_int - location[2]
+        if facing_direction.py_int == 2:  # north
+            if dz == 0 and dx == 1:
+                return {"connection": "right"}
+        elif facing_direction.py_int == 3:  # south
+            if dz == 0 and dx == -1:
+                return {"connection": "right"}
+        elif facing_direction.py_int == 4:  # west
+            if dx == 0 and dz == -1:
+                return {"connection": "right"}
+        elif facing_direction.py_int == 5:  # east
+            if dx == 0 and dz == 1:
+                return {"connection": "right"}
     return {}

PyMCTranslate/code_functions/bedrock_chest_connection_other_right_120.py

@@ -1,31 +1,35 @@
-def main(nbt, properties, location):
-    if (
-        nbt[0] == "compound"
-        and "pairlead" in nbt[1]
-        and nbt[1]["pairlead"][0] == "byte"
-        and nbt[1]["pairlead"][1] == 1
-        and "pairx" in nbt[1]
-        and nbt[1]["pairx"][0] == "int"
-        and "pairz" in nbt[1]
-        and nbt[1]["pairz"][0] == "int"
-        and "minecraft:cardinal_direction" in properties
-    ):
-        dx = nbt[1]["pairx"][1] - location[0]
-        dz = nbt[1]["pairz"][1] - location[2]
-        if properties["minecraft:cardinal_direction"].py_str == "north":  # north
-            if dz == 0:
-                if dx == 1:
-                    return {"connection": "right"}
-        elif properties["minecraft:cardinal_direction"].py_str == "south":  # south
-            if dz == 0:
-                if dx == -1:
-                    return {"connection": "right"}
-        elif properties["minecraft:cardinal_direction"].py_str == "west":  # west
-            if dx == 0:
-                if dz == -1:
-                    return {"connection": "right"}
-        elif properties["minecraft:cardinal_direction"].py_str == "east":  # east
-            if dx == 0:
-                if dz == 1:
-                    return {"connection": "right"}
+from typing import Union
+
+from amulet_nbt import CompoundTag, ByteTag, ShortTag, IntTag, LongTag, StringTag
+
+
+def main(nbt, properties: dict[str, Union[ByteTag, ShortTag, IntTag, LongTag, StringTag]], location: tuple[int, int, int]) -> dict[str, str]:
+    if not isinstance(nbt, CompoundTag):
+        return {}
+    facing_direction = properties.get("minecraft:cardinal_direction")
+    if not isinstance(facing_direction, StringTag):
+        return {}
+    pairlead = nbt.get("pairlead")
+    if isinstance(pairlead, ByteTag) and pairlead.py_int == 1:
+        pair_x = nbt.get("pairx")
+        if not isinstance(pair_x, IntTag):
+            return {}
+        pair_z = nbt.get("pairz")
+        if not isinstance(pair_z, IntTag):
+            return {}
+
+        dx = pair_x.py_int - location[0]
+        dz = pair_z.py_int - location[2]
+        if facing_direction.py_str == "north":  # north
+            if dz == 0 and dx == 1:
+                return {"connection": "right"}
+        elif facing_direction.py_str == "south":  # south
+            if dz == 0 and dx == -1:
+                return {"connection": "right"}
+        elif facing_direction.py_str == "west":  # west
+            if dx == 0 and dz == -1:
+                return {"connection": "right"}
+        elif facing_direction.py_str == "east":  # east
+            if dx == 0 and dz == 1:
+                return {"connection": "right"}
     return {}