Add TreeData#addRootItems, getRootItems

2 files changed, 80 insertions, 0 deletions

diff --git a/server/src/main/java/com/vaadin/data/TreeData.java b/server/src/main/java/com/vaadin/data/TreeData.java
index 2311a71984..7f6cc095bc 100644
--- a/server/src/main/java/com/vaadin/data/TreeData.java
+++ b/server/src/main/java/com/vaadin/data/TreeData.java
@@ -97,6 +97,60 @@ public class TreeData<T> implements Serializable {
     }
 
     /**
+     * Adds the items as root items to this structure.
+     *
+     * @param items
+     *            the items to add
+     * @return this
+     *
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are {code null}
+     */
+    public TreeData<T> addRootItems(T... items) {
+        addItems(null, items);
+        return this;
+    }
+
+    /**
+     * Adds the items of the given collection as root items to this structure.
+     *
+     * @param items
+     *            the collection of items to add
+     * @return this
+     *
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are {code null}
+     */
+    public TreeData<T> addRootItems(Collection<T> items) {
+        addItems(null, items);
+        return this;
+    }
+
+    /**
+     * Adds the items of the given stream as root items to this structure.
+     *
+     * @param items
+     *            the stream of root items to add
+     * @return this
+     *
+     * @throws IllegalArgumentException
+     *             if any of the given items have already been added to this
+     *             structure
+     * @throws NullPointerException
+     *             if any of the items are {code null}
+     */
+    public TreeData<T> addRootItems(Stream<T> items) {
+        addItems(null, items);
+        return this;
+    }
+
+    /**
      * Adds a data item as a child of {@code parent}. Call with {@code null} as
      * parent to add a root level item. The given parent item must already exist
      * in this structure, and an item can only be added to this structure once.
@@ -284,6 +338,15 @@ public class TreeData<T> implements Serializable {
     }
 
     /**
+     * Gets the root items of this structure.
+     *
+     * @return the root items of this structure
+     */
+    public List<T> getRootItems() {
+        return getChildren(null);
+    }
+
+    /**
      * Get the immediate child items for the given item.
      *
      * @param item
diff --git a/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java b/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
index f0b2d5f81f..186ed14e8a 100644
--- a/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/TreeDataProviderTest.java
@@ -81,6 +81,23 @@ public class TreeDataProviderTest
     }
 
     @Test
+    public void treeData_root_items() {
+        TreeData<String> data = new TreeData<>();
+        TreeData<String> dataVarargs = new TreeData<>();
+        TreeData<String> dataCollection = new TreeData<>();
+        TreeData<String> dataStream = new TreeData<>();
+
+        data.addItems(null, "a", "b", "c");
+        dataVarargs.addRootItems("a", "b", "c");
+        dataCollection.addRootItems(Arrays.asList("a", "b", "c"));
+        dataStream.addRootItems(Arrays.asList("a", "b", "c").stream());
+
+        Assert.assertEquals(data.getRootItems(), dataVarargs.getRootItems());
+        Assert.assertEquals(data.getRootItems(), dataCollection.getRootItems());
+        Assert.assertEquals(data.getRootItems(), dataStream.getRootItems());
+    }
+
+    @Test
     public void populate_treeData_with_child_item_provider() {
         TreeData<String> stringData = new TreeData<>();
         List<String> rootItems = Arrays.asList("a", "b", "c");