Simplify usage of selection.join with transitions

README.md

@@ -571,7 +571,7 @@ svg.selectAll("circle")
 
 The selections returned by the *enter* and *update* functions are merged and then returned by *selection*.join.
 
-You also animate enter, update and exit by creating transitions inside the *enter*, *update* and *exit* functions. To avoid breaking the method chain, use *selection*.call to create transitions, or return an undefined enter or update selection to prevent merging: the return value of the *enter* and *update* functions specifies the two selections to merge and return by *selection*.join.
+You can animate enter, update and exit by creating transitions inside the *enter*, *update* and *exit* functions. If the *enter* and *update* functions return transitions, their underlying selections are merged and then returned by *selection*.join. The return value of the *exit* function is not used.
 
 For more, see the [*selection*.join notebook](https://observablehq.com/@d3/selection-join).
 

src/selection/join.js

@@ -1,7 +1,15 @@
 export default function(onenter, onupdate, onexit) {
   var enter = this.enter(), update = this, exit = this.exit();
-  enter = typeof onenter === "function" ? onenter(enter) : enter.append(onenter + "");
-  if (onupdate != null) update = onupdate(update);
+  if (typeof onenter === "function") {
+    enter = onenter(enter);
+    if (enter) enter = enter.selection();
+  } else {
+    enter = enter.append(onenter + "");
+  }
+  if (onupdate != null) {
+    update = onupdate(update);
+    if (update) update = update.selection();
+  }
   if (onexit == null) exit.remove(); else onexit(exit);
   return enter && update ? enter.merge(update).order() : update;
 }

test/selection/join-test.js

@@ -39,3 +39,21 @@ tape("selection.join(…) reorders nodes to match the data", function(test) {
   test.equal(document.body.innerHTML, "<p>0</p><p>3</p><p>1</p><p>2</p><p>4</p>");
   test.end();
 });
+
+function mockTransition(selection){
+  return {
+    selection: function() { return selection; }
+  }
+}
+
+tape("selection.join(enter, update, exit) allows callbacks to return a transition", function(test) {
+  var document = jsdom("<p>1</p><p>2</p>"),
+      p = d3.select(document.body).selectAll("p").datum(function() { return this.textContent; });
+  p = p.data([1, 3], d => d).join(
+    enter => mockTransition(enter.append("p").attr("class", "enter").text(d => d)),
+    update => mockTransition(update.attr("class", "update")),
+    exit => mockTransition(exit.attr("class", "exit"))
+  );
+  test.equal(document.body.innerHTML, "<p class=\"update\">1</p><p class=\"exit\">2</p><p class=\"enter\">3</p>");
+  test.end();
+});