Improve method U.set(map, path, value).

Author: javadev

src/main/java/com/github/underscore/lodash/U.java

@@ -1497,28 +1497,42 @@ public static List<String> stringToPath(final String string) {
     }
 
     @SuppressWarnings("unchecked")
-    private static <T> T baseGet(final Map<String, Object> object, final String path) {
+    private static <T> T baseGetAndSet(final Map<String, Object> object, final String path,
+        final Optional<Object> value) {
         final List<String> paths = stringToPath(path);
         int index = 0;
         final int length = paths.size();
 
         Object localObject = object;
+        Object savedLocalObject = null;
+        String savedPath = null;
         while (localObject != null && index < length) {
             if (localObject instanceof Map) {
                 Map.Entry mapEntry = getMapEntry((Map) localObject);
                 if (mapEntry != null && "#item".equals(mapEntry.getKey())) {
                     localObject = mapEntry.getValue();
                     continue;
                 }
+                savedLocalObject = localObject;
+                savedPath = paths.get(index);
                 localObject = ((Map) localObject).get(paths.get(index));
             } else if (localObject instanceof List) {
+                savedLocalObject = localObject;
+                savedPath = paths.get(index);
                 localObject = ((List) localObject).get(Integer.parseInt(paths.get(index)));
             } else {
                 break;
             }
             index += 1;
         }
         if (index > 0 && index == length) {
+            if (value.isPresent()) {
+                if (savedLocalObject instanceof Map) {
+                    ((Map) savedLocalObject).put(savedPath, value.get());
+                } else {
+                    ((List) savedLocalObject).set(Integer.parseInt(savedPath), value.get());
+                }
+            }
             return (T) localObject;
         }
         return null;
@@ -1529,50 +1543,11 @@ private static Map.Entry getMapEntry(Map map) {
     }
 
     public static <T> T get(final Map<String, Object> object, final String path) {
-        return baseGet(object, path);
-    }
-
-    @SuppressWarnings("unchecked")
-    private static Map<String, Object> baseSet(final Map<String, Object> object,
-            final String path, final Object value) {
-        final List<String> paths = stringToPath(path);
-        int index = 0;
-        final int length = paths.size();
-
-        Object savedLocalObject = object;
-        String savedPath = paths.get(0);
-        Object localObject = object;
-        while (localObject != null && index < length) {
-            if (localObject instanceof Map) {
-                Map.Entry mapEntry = getMapEntry((Map) localObject);
-                if (mapEntry != null && "#item".equals(mapEntry.getKey())) {
-                    localObject = mapEntry.getValue();
-                    continue;
-                }
-                savedLocalObject = localObject;
-                savedPath = paths.get(index);
-                localObject = ((Map) localObject).get(paths.get(index));
-            } else if (localObject instanceof List) {
-                savedLocalObject = localObject;
-                savedPath = paths.get(index);
-                localObject = ((List) localObject).get(Integer.parseInt(paths.get(index)));
-            } else {
-                break;
-            }
-            index += 1;
-        }
-        if (index > 0 && index == length) {
-            if (savedLocalObject instanceof Map) {
-                ((Map) savedLocalObject).put(savedPath, value);
-            } else {
-                ((List) savedLocalObject).set(Integer.parseInt(savedPath), value);
-            }
-        }
-        return object;
+        return baseGetAndSet(object, path, Optional.absent());
     }
 
-    public static Map<String, Object> set(final Map<String, Object> object, final String path, Object value) {
-        return baseSet(object, path, value);
+    public static <T> T set(final Map<String, Object> object, final String path, Object value) {
+        return baseGetAndSet(object, path, Optional.of(value));
     }
 
     public static class FetchResponse {

src/test/java/com/github/underscore/lodash/LodashTest.java

@@ -442,22 +442,22 @@ public void get() {
     @SuppressWarnings("unchecked")
     @Test
     public void set() {
-        assertEquals("{a=[{b={c=e}}]}", U.<String>set(
+        assertEquals("d", U.<String>set(
             (Map<String, Object>) U.fromJson("{\"a\":[{\"b\":{\"c\":\"d\"}}]}"), "a[0].b.c", "e").toString());
-        assertEquals("{a=[e]}", U.<String>set(
+        assertEquals("{b={c=d}}", U.set(
             (Map<String, Object>) U.fromJson("{\"a\":[{\"b\":{\"c\":\"d\"}}]}"), "a[0]", "e").toString());
         Map<String, Object> map = U.newLinkedHashMap();
         Map<String, Object> map2 = U.newLinkedHashMap();
         Map<String, Object> map3 = U.newLinkedHashMap();
         map.put("a", map2);
         map2.put("#item", map3);
         map3.put("b", "c");
-        assertEquals("{a={#item={b=b}}}", U.<String>set(map, "a.b", "b").toString());
+        assertEquals("c", U.<String>set(map, "a.b", "b").toString());
         assertNull(U.<String>set((Map<String, Object>) null, "a", "b"));
-        assertEquals("{}", U.<String>set(U.<String, Object>newLinkedHashMap(), "a.b", "b").toString());
+        assertNull(U.<String>set(U.<String, Object>newLinkedHashMap(), "a.b", "b"));
         Map<String, Object> map4 = U.newLinkedHashMap();
         map4.put("a", "b");
-        assertEquals("{a=b}", U.<String>set(map4, "a.b", "b").toString());
+        assertNull(U.<String>set(map4, "a.b", "b"));
     }
 
     @Test