From b547df5cfc55aad6177b61f37c9ef88f1877b884 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philippe=20Rivi=C3=A8re?= <fil@rezo.net>
Date: Wed, 19 Jul 2023 09:55:02 +0200
Subject: [PATCH] imputeLeaf closes #206

---
 README.md             |  8 ++++++++
 src/stratify.js       | 19 +++++++++++++++++++
 test/stratify-test.js | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/README.md b/README.md
index 69e4f7e8..6411fe7c 100644
--- a/README.md
+++ b/README.md
@@ -334,6 +334,14 @@ If *path* is specified, sets the path accessor to the given function and returns
 d3.stratify().path(d => d)(["a/b", "a/c"]); // nodes with id "/a", "/a/b", "/a/c"
 ```
 
+<a name="stratify_imputeleaf" href="#stratify_imputeleaf">#</a> <i>stratify</i>.<b>imputeLeaf</b>([<i>impute</i>]) · [Source](https://github.com/d3/d3-hierarchy/blob/main/src/stratify.js), [Examples](https://observablehq.com/@d3/d3-stratify)
+
+If *impute* is specified as a boolean, sets the imputeLeaf option to the given value and returns this stratify operator. Otherwise, returns the current impute option, which defaults to false. If the option is true, a node that has children and a value will be represented as a branching node with no data, with a child node bearing the data. The child node’s id will be the parent id suffixed with a slash.
+
+```js
+d3.stratify().path(d => d).imputeLeaf(true)(["a", "a/b"]); // nodes with id "/a", "/a/", "/a/b"
+```
+
 ### Cluster
 
 [<img alt="Dendrogram" src="https://raw.githubusercontent.com/d3/d3-hierarchy/main/img/cluster.png">](https://observablehq.com/@d3/cluster-dendrogram)
diff --git a/src/stratify.js b/src/stratify.js
index d113100d..b0fc3e00 100644
--- a/src/stratify.js
+++ b/src/stratify.js
@@ -16,6 +16,7 @@ function defaultParentId(d) {
 export default function() {
   var id = defaultId,
       parentId = defaultParentId,
+      imputeLeaf = false,
       path;
 
   function stratify(data) {
@@ -94,6 +95,20 @@ export default function() {
     root.parent = null;
     if (n > 0) throw new Error("cycle");
 
+    if (imputeLeaf) {
+      root.each(node => {
+        if (node.children && node.children.length && node.data != null) {
+          const child = new Node(node.data);
+          node.data = null;
+          child.depth = node.depth + 1;
+          child.height = 0;
+          child.parent = node;
+          child.id = node.id + "/";
+          node.children.unshift(child);
+        }
+      });
+    }
+
     return root;
   }
 
@@ -109,6 +124,10 @@ export default function() {
     return arguments.length ? (path = optional(x), stratify) : path;
   };
 
+  stratify.imputeLeaf = function (x) {
+    return arguments.length ? (imputeLeaf = !!x, stratify) : this.imputeLeaf;
+  }
+
   return stratify;
 }
 
diff --git a/test/stratify-test.js b/test/stratify-test.js
index e6029f73..4177a9a9 100644
--- a/test/stratify-test.js
+++ b/test/stratify-test.js
@@ -227,6 +227,41 @@ it("stratify(data) does not treat a falsy but non-empty parentId as the root", (
   });
 });
 
+it("stratify(data).imputeLeaf(true) imputes leaves when branch nodes have a value", () => {
+  const s = stratify().path(d => d.path).imputeLeaf(true);
+  const root = s([
+    {path: "/gallery", visits: 4},
+    {path: "/gallery/a", visits: 2},
+    {path: "/gallery/b", visits: 1}
+  ]);
+  assert.deepStrictEqual(noparent(root), {
+    "id": "/gallery",
+    "depth": 0,
+    "height": 1,
+    "data": null,
+    "children": [
+      {
+        "id": "/gallery/",
+        "depth": 1,
+        "height": 0,
+        "data": {"path": "/gallery", "visits": 4},
+      },
+      {
+        "id": "/gallery/a",
+        "depth": 1,
+        "height": 0,
+        "data": {"path": "/gallery/a", "visits": 2}
+      },
+      {
+        "id": "/gallery/b",
+        "depth": 1,
+        "height": 0,
+        "data": {"path": "/gallery/b", "visits": 1}
+      }
+    ]
+  });
+});
+
 it("stratify(data) throws an error if the data does not have a single root", () => {
   const s = stratify();
   assert.throws(() => { s([{id: "a"}, {id: "b"}]); }, /\bmultiple roots\b/);