diff --git a/src/jsoneditor/components/TextMode.js b/src/jsoneditor/components/TextMode.js
index 3e810c2..a696021 100644
--- a/src/jsoneditor/components/TextMode.js
+++ b/src/jsoneditor/components/TextMode.js
@@ -62,7 +62,7 @@ export default class TextMode extends Component {
     this.applyProps(nextProps, this.props)
   }
 
-  // TODO: create some sort of watcher structure for these props? Is there a Reactpattern for that?
+  // TODO: create some sort of watcher structure for these props? Is there a React pattern for that?
   applyProps (nextProps, currentProps) {
     // Apply text
     if (nextProps.text !== currentProps.text) {
@@ -87,6 +87,7 @@ export default class TextMode extends Component {
       this.findKeyBinding = createFindKeyBinding(keyBindings)
     }
 
+    // TODO: apply patch
     // TODO: apply patchText
   }
 
diff --git a/src/jsoneditor/components/TreeMode.js b/src/jsoneditor/components/TreeMode.js
index 675d485..e2dfb60 100644
--- a/src/jsoneditor/components/TreeMode.js
+++ b/src/jsoneditor/components/TreeMode.js
@@ -156,9 +156,9 @@ export default class TreeMode extends PureComponent {
     // Apply json
     if (nextProps.json !== this.state.json) {
       // FIXME: merge meta data from existing eson
-      const expandCallback = this.props.expand || TreeMode.expandRoot
+      const callback = this.props.expand || TreeMode.expandRoot
       const json = nextProps.json
-      const eson = expand(jsonToEson(json), expandCallback)
+      const eson = expand(jsonToEson(json), callback)
 
       this.setState({
         json,
@@ -571,7 +571,7 @@ export default class TreeMode extends PureComponent {
     if (recurse) {
       this.setState({
         eson: updateIn(this.state.eson, path, function (child) {
-          return expand(child, (path) => true, expanded)
+          return expand(child, (path) => expanded)
         })
       })
     }
@@ -588,18 +588,14 @@ export default class TreeMode extends PureComponent {
   }
 
   handleExpandAll = () => {
-    const expanded = true
-
     this.setState({
-      eson: expand(this.state.eson, TreeMode.expandAll, expanded)
+      eson: expand(this.state.eson, TreeMode.expandAll)
     })
   }
 
   handleCollapseAll = () => {
-    const expanded = false
-
     this.setState({
-      eson: expand(this.state.eson, TreeMode.expandAll, expanded)
+      eson: expand(this.state.eson, TreeMode.collapseAll)
     })
   }
 
@@ -1008,9 +1004,18 @@ export default class TreeMode extends PureComponent {
    * @param {Path | function (path: Path) : boolean} callback
    */
   expand (callback) {
-    this.setState({
-      eson: expand(this.state.eson, callback, true)
-    })
+    if (Array.isArray(callback)) {
+      this.setState({
+        eson: expandPath(this.state.eson, callback, true)
+      })
+    }
+    else { // callback is a function
+      this.setState({
+        eson: expand(this.state.eson, (path) => {
+          return callback(path) === true ? true : undefined
+        })
+      })
+    }
   }
 
   /**
@@ -1018,9 +1023,18 @@ export default class TreeMode extends PureComponent {
    * @param {Path | function (path: Path) : boolean} callback
    */
   collapse (callback) {
-    this.setState({
-      eson: expand(this.state.eson, callback, false)
-    })
+    if (Array.isArray(callback)) {
+      this.setState({
+        eson: expandPath(this.state.eson, callback, true)
+      })
+    }
+    else { // callback is a function
+      this.setState({
+        eson: expand(this.state.eson, (path) => {
+          return callback(path) === true ? false : undefined
+        })
+      })
+    }
   }
 
   /**
@@ -1074,6 +1088,16 @@ export default class TreeMode extends PureComponent {
   static expandAll (path) {
     return true
   }
+
+  /**
+   * Callback function to collapse all nodes
+   *
+   * @param {Array.<string>} path
+   * @return {boolean}
+   */
+  static collapseAll (path) {
+    return false
+  }
 }
 
 // TODO: describe PropTypes
diff --git a/src/jsoneditor/eson.js b/src/jsoneditor/eson.js
index 356401c..0d3a5d5 100644
--- a/src/jsoneditor/eson.js
+++ b/src/jsoneditor/eson.js
@@ -138,19 +138,23 @@ export function updatePaths(eson, path = []) {
 /**
  * Expand or collapse all items matching a filter callback
  * @param {ESON} eson
- * @param {function(Path) : boolean | Path} filterCallback
- *              When a path, the object/array at this path will be expanded/collapsed
- *              When a function, all objects and arrays for which callback
- *              returns true will be expanded/collapsed
- * @param {boolean} [expanded=true]  New expanded state: true to expand, false to collapse
+ * @param {function(Path) : boolean | undefined} callback
+ *              All objects and arrays for which callback returns true will be expanded
+ *              All objects and arrays for which callback returns false will be collapsed
+ *              All objects and arrays for which callback returns undefined will be left as is
  * @return {ESON}
  */
-export function expand (eson, filterCallback, expanded = true) {
-  // TODO: adjust expand to have a filterCallback which can return true, false, or undefined. In the latter case, the expanded state is left as is.
+export function expand (eson, callback) {
   return transform(eson, function (value, path) {
-    return ((value[META].type === 'Array' || value[META].type === 'Object') && filterCallback(path))
-        ? expandOne(value, [], expanded)
-        : value
+    if (value[META].type === 'Array' || value[META].type === 'Object') {
+      const expanded = callback(path)
+      return (typeof expanded === 'boolean')
+          ? expandOne(value, [], expanded) // adjust expanded state
+          : value  // leave as is when returned value is null or undefined
+    }
+    else {
+      return value
+    }
   })
 }
 
diff --git a/src/jsoneditor/eson.test.js b/src/jsoneditor/eson.test.js
index 57525e6..f3ea5ba 100644
--- a/src/jsoneditor/eson.test.js
+++ b/src/jsoneditor/eson.test.js
@@ -112,15 +112,16 @@ test('expand a callback', () => {
     "bool": false
   })
 
-  function filterCallback (path) {
-    return path.length >= 1
+  function callback (path) {
+    return (path.length >= 1)
+        ? false     // collapse
+        : undefined // leave untouched
   }
-  const expandedValue = false
-  const collapsed = expand(eson, filterCallback, expandedValue)
+  const collapsed = expand(eson, callback)
   expect(collapsed[META].expanded).toEqual(undefined)
-  expect(collapsed.obj[META].expanded).toEqual(expandedValue)
-  expect(collapsed.obj.arr[META].expanded).toEqual(expandedValue)
-  expect(collapsed.obj.arr[2][META].expanded).toEqual(expandedValue)
+  expect(collapsed.obj[META].expanded).toEqual(false)
+  expect(collapsed.obj.arr[META].expanded).toEqual(false)
+  expect(collapsed.obj.arr[2][META].expanded).toEqual(false)
 
   let orig = collapsed
   orig = deleteIn(orig, ['obj'].concat([META, 'expanded']))
@@ -132,10 +133,9 @@ test('expand a callback', () => {
 test('expand a callback should not change the object when nothing happens', () => {
   const eson = jsonToEson({a: [1,2,3], b: {c: 4}})
   function callback (path) {
-    return false
+    return undefined
   }
-  const expanded = false
-  const collapsed = expand(eson, callback, expanded)
+  const collapsed = expand(eson, callback)
 
   expect(collapsed).toBe(eson)
 })