fixed a bug where topEdge was not set in _get_positions_and_total_wid…

…ths (#45)
Pierre-Sassoulas · Jun 14, 2024 · 4550755 · 4550755
1 parent b06fc3f
commit 4550755
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/pysankey/sankey/__init__.py b/pysankey/sankey/__init__.py
@@ -493,7 +493,7 @@ def _get_positions_and_total_widths(
             weighted_sum = 0.02 * df[side + "Weight"].sum()
             label_widths["bottom"] = bottom_width + weighted_sum
             label_widths["top"] = label_widths["bottom"] + label_widths[side]
-            topEdge = label_widths["top"]
+        topEdge = label_widths["top"]
         widths[label] = label_widths
         LOGGER.debug("%s position of '%s' : %s", side, label, label_widths)
     # pylint: disable-next=possibly-used-before-assignment

diff --git a/tests/test_sankey.py b/tests/test_sankey.py
@@ -8,3 +8,8 @@ class TestSankey(TestFruit):
     def test_right_color(self) -> None:
         ax = sankey(self.data["true"], self.data["predicted"], rightColor=True)
         self.assertIsInstance(ax, plt.Axes)
+
+    def test_single(self) -> None:
+        source = [1]
+        target = [2]
+        sankey(source, target, rightColor=True)