Add support for entities with .schematic files.

src/main/java/com/sk89q/worldedit/extent/clipboard/io/SchematicReader.java

@@ -21,6 +21,8 @@ package com.sk89q.worldedit.extent.clipboard.io;
 
 import com.sk89q.jnbt.ByteArrayTag;
 import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.jnbt.DoubleTag;
+import com.sk89q.jnbt.FloatTag;
 import com.sk89q.jnbt.IntTag;
 import com.sk89q.jnbt.ListTag;
 import com.sk89q.jnbt.NBTInputStream;
@@ -31,12 +33,15 @@ import com.sk89q.worldedit.BlockVector;
 import com.sk89q.worldedit.Vector;
 import com.sk89q.worldedit.WorldEditException;
 import com.sk89q.worldedit.blocks.BaseBlock;
+import com.sk89q.worldedit.entity.BaseEntity;
 import com.sk89q.worldedit.extent.clipboard.BlockArrayClipboard;
 import com.sk89q.worldedit.extent.clipboard.Clipboard;
 import com.sk89q.worldedit.regions.CuboidRegion;
 import com.sk89q.worldedit.regions.Region;
+import com.sk89q.worldedit.util.Location;
 import com.sk89q.worldedit.world.registry.WorldData;
 
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.List;
@@ -79,29 +84,32 @@ public class SchematicReader implements ClipboardReader {
         }
 
         // Check type of Schematic
-        String materials = getChildTag(schematic, "Materials", StringTag.class).getValue();
+        String materials = requireTag(schematic, "Materials", StringTag.class).getValue();
         if (!materials.equals("Alpha")) {
             throw new IOException("Schematic file is not an Alpha schematic");
         }
 
-        // Parse origin and region from WEOrigin and WEOffset
+        // ====================================================================
+        // Metadata
+        // ====================================================================
+
         Vector origin;
         Region region;
 
         // Get information
-        short width = getChildTag(schematic, "Width", ShortTag.class).getValue();
+        short width = requireTag(schematic, "Width", ShortTag.class).getValue();
-        short height = getChildTag(schematic, "Height", ShortTag.class).getValue();
+        short height = requireTag(schematic, "Height", ShortTag.class).getValue();
-        short length = getChildTag(schematic, "Length", ShortTag.class).getValue();
+        short length = requireTag(schematic, "Length", ShortTag.class).getValue();
 
         try {
-            int originX = getChildTag(schematic, "WEOriginX", IntTag.class).getValue();
+            int originX = requireTag(schematic, "WEOriginX", IntTag.class).getValue();
-            int originY = getChildTag(schematic, "WEOriginY", IntTag.class).getValue();
+            int originY = requireTag(schematic, "WEOriginY", IntTag.class).getValue();
-            int originZ = getChildTag(schematic, "WEOriginZ", IntTag.class).getValue();
+            int originZ = requireTag(schematic, "WEOriginZ", IntTag.class).getValue();
             Vector min = new Vector(originX, originY, originZ);
 
-            int offsetX = getChildTag(schematic, "WEOffsetX", IntTag.class).getValue();
+            int offsetX = requireTag(schematic, "WEOffsetX", IntTag.class).getValue();
-            int offsetY = getChildTag(schematic, "WEOffsetY", IntTag.class).getValue();
+            int offsetY = requireTag(schematic, "WEOffsetY", IntTag.class).getValue();
-            int offsetZ = getChildTag(schematic, "WEOffsetZ", IntTag.class).getValue();
+            int offsetZ = requireTag(schematic, "WEOffsetZ", IntTag.class).getValue();
             Vector offset = new Vector(offsetX, offsetY, offsetZ);
 
             origin = min.subtract(offset);
@@ -111,16 +119,20 @@ public class SchematicReader implements ClipboardReader {
             region = new CuboidRegion(origin, origin.add(width, height, length).subtract(Vector.ONE));
         }
 
+        // ====================================================================
+        // Blocks
+        // ====================================================================
+
         // Get blocks
-        byte[] blockId = getChildTag(schematic, "Blocks", ByteArrayTag.class).getValue();
+        byte[] blockId = requireTag(schematic, "Blocks", ByteArrayTag.class).getValue();
-        byte[] blockData = getChildTag(schematic, "Data", ByteArrayTag.class).getValue();
+        byte[] blockData = requireTag(schematic, "Data", ByteArrayTag.class).getValue();
         byte[] addId = new byte[0];
         short[] blocks = new short[blockId.length]; // Have to later combine IDs
 
         // We support 4096 block IDs using the same method as vanilla Minecraft, where
         // the highest 4 bits are stored in a separate byte array.
         if (schematic.containsKey("AddBlocks")) {
-            addId = getChildTag(schematic, "AddBlocks", ByteArrayTag.class).getValue();
+            addId = requireTag(schematic, "AddBlocks", ByteArrayTag.class).getValue();
         }
 
         // Combine the AddBlocks data with the first 8-bit block ID
@@ -137,7 +149,7 @@ public class SchematicReader implements ClipboardReader {
         }
 
         // Need to pull out tile entities
-        List<Tag> tileEntities = getChildTag(schematic, "TileEntities", ListTag.class).getValue();
+        List<Tag> tileEntities = requireTag(schematic, "TileEntities", ListTag.class).getValue();
         Map<BlockVector, Map<String, Tag>> tileEntitiesMap = new HashMap<BlockVector, Map<String, Tag>>();
 
         for (Tag tag : tileEntities) {
@@ -208,19 +220,64 @@ public class SchematicReader implements ClipboardReader {
             }
         }
 
+        // ====================================================================
+        // Entities
+        // ====================================================================
+
+        try {
+            List<Tag> entityTags = requireTag(schematic, "Entities", ListTag.class).getValue();
+
+            for (Tag tag : entityTags) {
+                if (tag instanceof CompoundTag) {
+                    CompoundTag compound = (CompoundTag) tag;
+                    StringTag idTag = getTag(compound, StringTag.class, "id");
+                    Vector position = getVector(getTag(compound, ListTag.class, "Pos"));
+
+                    if (idTag != null & position != null) {
+                        Location location = readRotation(getTag(compound, ListTag.class, "Rotation"), new Location(clipboard, position));
+                        BaseEntity state = new BaseEntity(idTag.getValue(), compound);
+                        clipboard.createEntity(location, state);
+                    }
+                }
+            }
+        } catch (IOException ignored) { // No entities? No problem
+        }
+
         return clipboard;
     }
 
-    /**
+    @Nullable
-     * Get child tag of a NBT structure.
+    private static Vector getVector(@Nullable ListTag tag) {
-     *
+        if (tag != null) {
-     * @param items The parent tag map
+            List<Tag> tags = tag.getValue();
-     * @param key The name of the tag to get
+
-     * @param expected The expected type of the tag
+            if (tags.size() == 3 && tags.get(0) instanceof DoubleTag) {
-     * @return child tag casted to the expected type
+                double x = ((DoubleTag) tags.get(0)).getValue();
-     * @throws IOException if the tag does not exist or the tag is not of the expected type
+                double y = ((DoubleTag) tags.get(1)).getValue();
-     */
+                double z = ((DoubleTag) tags.get(2)).getValue();
-    private static <T extends Tag> T getChildTag(Map<String, Tag> items, String key, Class<T> expected) throws IOException {
+                return new Vector(x, y, z);
+            }
+        }
+
+        return null;
+    }
+
+    @Nullable
+    private static Location readRotation(@Nullable ListTag tag, Location location) {
+        if (tag != null) {
+            List<Tag> tags = tag.getValue();
+
+            if (tags.size() == 2 && tags.get(0) instanceof FloatTag) {
+                float yaw = ((FloatTag) tags.get(0)).getValue();
+                float pitch = ((FloatTag) tags.get(1)).getValue();
+                location = location.setDirection(yaw, pitch);
+            }
+        }
+
+        return location;
+    }
+
+    private static <T extends Tag> T requireTag(Map<String, Tag> items, String key, Class<T> expected) throws IOException {
         if (!items.containsKey(key)) {
             throw new IOException("Schematic file is missing a \"" + key + "\" tag");
         }
@@ -233,4 +290,20 @@ public class SchematicReader implements ClipboardReader {
         return expected.cast(tag);
     }
 
+    @Nullable
+    private static <T extends Tag> T getTag(CompoundTag tag, Class<T> expected, String key) {
+        Map<String, Tag> items = tag.getValue();
+
+        if (!items.containsKey(key)) {
+            return null;
+        }
+
+        Tag test = items.get(key);
+        if (!expected.isInstance(test)) {
+            return null;
+        }
+
+        return expected.cast(test);
+    }
+
 }

src/main/java/com/sk89q/worldedit/extent/clipboard/io/SchematicWriter.java

@@ -21,6 +21,8 @@ package com.sk89q.worldedit.extent.clipboard.io;
 
 import com.sk89q.jnbt.ByteArrayTag;
 import com.sk89q.jnbt.CompoundTag;
+import com.sk89q.jnbt.DoubleTag;
+import com.sk89q.jnbt.FloatTag;
 import com.sk89q.jnbt.IntTag;
 import com.sk89q.jnbt.ListTag;
 import com.sk89q.jnbt.NBTOutputStream;
@@ -29,19 +31,23 @@ import com.sk89q.jnbt.StringTag;
 import com.sk89q.jnbt.Tag;
 import com.sk89q.worldedit.Vector;
 import com.sk89q.worldedit.blocks.BaseBlock;
+import com.sk89q.worldedit.entity.BaseEntity;
+import com.sk89q.worldedit.entity.Entity;
 import com.sk89q.worldedit.extent.clipboard.Clipboard;
 import com.sk89q.worldedit.regions.Region;
+import com.sk89q.worldedit.util.Location;
 import com.sk89q.worldedit.world.registry.WorldData;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
- /**
+/**
  * Writes schematic files based that are compatible with MCEdit and other editors.
  */
 public class SchematicWriter implements ClipboardWriter {
@@ -79,6 +85,10 @@ public class SchematicWriter implements ClipboardWriter {
             throw new IllegalArgumentException("Length of region too large for a .schematic");
         }
 
+        // ====================================================================
+        // Metadata
+        // ====================================================================
+
         HashMap<String, Tag> schematic = new HashMap<String, Tag>();
         schematic.put("Width", new ShortTag("Width", (short) width));
         schematic.put("Length", new ShortTag("Length", (short) length));
@@ -91,11 +101,14 @@ public class SchematicWriter implements ClipboardWriter {
         schematic.put("WEOffsetY", new IntTag("WEOffsetY", offset.getBlockY()));
         schematic.put("WEOffsetZ", new IntTag("WEOffsetZ", offset.getBlockZ()));
 
-        // Copy
+        // ====================================================================
+        // Block handling
+        // ====================================================================
+
         byte[] blocks = new byte[width * height * length];
         byte[] addBlocks = null;
         byte[] blockData = new byte[width * height * length];
-        ArrayList<Tag> tileEntities = new ArrayList<Tag>();
+        List<Tag> tileEntities = new ArrayList<Tag>();
 
         for (Vector point : region) {
             Vector relative = point.subtract(min);
@@ -140,20 +153,66 @@ public class SchematicWriter implements ClipboardWriter {
 
         schematic.put("Blocks", new ByteArrayTag("Blocks", blocks));
         schematic.put("Data", new ByteArrayTag("Data", blockData));
-        schematic.put("Entities", new ListTag("Entities", CompoundTag.class, new ArrayList<Tag>()));
         schematic.put("TileEntities", new ListTag("TileEntities", CompoundTag.class, tileEntities));
 
         if (addBlocks != null) {
             schematic.put("AddBlocks", new ByteArrayTag("AddBlocks", addBlocks));
         }
 
-        // Build and output
+        // ====================================================================
+        // Entities
+        // ====================================================================
+
+        List<Tag> entities = new ArrayList<Tag>();
+        for (Entity entity : clipboard.getEntities()) {
+            BaseEntity state = entity.getState();
+
+            if (state != null) {
+                Map<String, Tag> values = new HashMap<String, Tag>();
+
+                // Put NBT provided data
+                CompoundTag rawTag = state.getNbtData();
+                if (rawTag != null) {
+                    values.putAll(rawTag.getValue());
+                }
+
+                // Store our location data, overwriting any
+                values.put("id", new StringTag("id", state.getTypeId()));
+                values.put("Pos", writeVector(entity.getLocation().toVector(), "Pos"));
+                values.put("Rotation", writeRotation(entity.getLocation(), "Rotation"));
+
+                CompoundTag entityTag = new CompoundTag("Entity", values);
+                entities.add(entityTag);
+            }
+        }
+
+        schematic.put("Entities", new ListTag("Entities", CompoundTag.class, entities));
+
+        // ====================================================================
+        // Output
+        // ====================================================================
+
         CompoundTag schematicTag = new CompoundTag("Schematic", schematic);
         outputStream.writeTag(schematicTag);
     }
 
+    private Tag writeVector(Vector vector, String name) {
+        List<DoubleTag> list = new ArrayList<DoubleTag>();
+        list.add(new DoubleTag("", vector.getX()));
+        list.add(new DoubleTag("", vector.getY()));
+        list.add(new DoubleTag("", vector.getZ()));
+        return new ListTag(name, DoubleTag.class, list);
+    }
+
+    private Tag writeRotation(Location location, String name) {
+        List<FloatTag> list = new ArrayList<FloatTag>();
+        list.add(new FloatTag("", location.getYaw()));
+        list.add(new FloatTag("", location.getPitch()));
+        return new ListTag(name, FloatTag.class, list);
+    }
+
     @Override
     public void close() throws IOException {
         outputStream.close();
     }
- }
+}