Transform polar segment endpoints

Author: iangow

plotnine/coords/coord_polar.py

@@ -130,19 +130,31 @@ def transform(
 
         if self.theta == "x":
             theta_col, r_col = "x", "y"
+            theta_end_col, r_end_col = "xend", "yend"
         else:
             theta_col, r_col = "y", "x"
+            theta_end_col, r_end_col = "yend", "xend"
 
         if theta_col not in data.columns or r_col not in data.columns:
             return data
 
         data = data.copy()
         data[theta_col] = self._to_radians(data[theta_col].to_numpy())
+        has_endpoints = theta_end_col in data.columns and r_end_col in data.columns
+        if has_endpoints:
+            data[theta_end_col] = self._to_radians(
+                data[theta_end_col].to_numpy()
+            )
 
         # PolarAxes always expects x = theta (radians) and y = r.
         # When theta = "y" we need to swap the columns.
         if self.theta == "y":
             data["x"], data["y"] = data["y"].copy(), data["x"].copy()
+            if has_endpoints:
+                data["xend"], data["yend"] = (
+                    data["yend"].copy(),
+                    data["xend"].copy(),
+                )
 
         return data
 

tests/test_coord_polar.py

@@ -0,0 +1,30 @@
+import numpy as np
+import pandas as pd
+
+from plotnine.coords.coord_polar import coord_polar
+
+
+def test_coord_polar_transforms_segment_endpoints_theta_x():
+    coord = coord_polar(theta="x")
+    coord.params = {"theta_range": (0, 10), "r_range": (0, 10)}
+    data = pd.DataFrame({"x": [0], "y": [1], "xend": [10], "yend": [2]})
+
+    out = coord.transform(data, None)
+
+    assert out.loc[0, "x"] == 0
+    assert out.loc[0, "y"] == 1
+    assert np.isclose(out.loc[0, "xend"], 2 * np.pi)
+    assert out.loc[0, "yend"] == 2
+
+
+def test_coord_polar_transforms_segment_endpoints_theta_y():
+    coord = coord_polar(theta="y")
+    coord.params = {"theta_range": (0, 10), "r_range": (0, 10)}
+    data = pd.DataFrame({"x": [1], "y": [0], "xend": [2], "yend": [10]})
+
+    out = coord.transform(data, None)
+
+    assert out.loc[0, "x"] == 0
+    assert out.loc[0, "y"] == 1
+    assert np.isclose(out.loc[0, "xend"], 2 * np.pi)
+    assert out.loc[0, "yend"] == 2