fix some of my messy previous change that broke theta_labels... (#146)

* fix some of my messy previous change that broke theta_labels...

* fix to support python2

* rebase and adapt to the new tests

Co-authored-by: Samuel <[email protected]>
Co-authored-by: Filipe Fernandes <[email protected]>

Author: 3 people

tests/test_windrose.py

@@ -1,10 +1,11 @@
 import matplotlib
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 from numpy.testing import assert_allclose
 from pandas.testing import assert_frame_equal
 
-from windrose import clean, clean_df, plot_windrose
+from windrose import DEFAULT_THETA_LABELS, WindroseAxes, clean, clean_df, plot_windrose
 
 matplotlib.use("Agg")  # noqa
 # Create wind speed and direction variables
@@ -52,3 +53,17 @@ def test_windrose_clean_df():
         }
     )
     assert_frame_equal(actual_df, expected_df)
+
+
+def test_theta_labels():
+    # Ensure default theta_labels are correct
+    ax = WindroseAxes.from_ax()
+    theta_labels = [t.get_text() for t in ax.get_xticklabels()]
+    assert theta_labels == DEFAULT_THETA_LABELS
+    plt.close()
+
+    # Ensure theta_labels are changed when specified
+    ax = WindroseAxes.from_ax(theta_labels=list("abcdefgh"))
+    theta_labels = [t.get_text() for t in ax.get_xticklabels()]
+    assert theta_labels == ["a", "b", "c", "d", "e", "f", "g", "h"]
+    plt.close()

windrose/__init__.py

@@ -3,6 +3,7 @@
 from matplotlib.projections import register_projection
 
 from .windrose import D_KIND_PLOT  # noqa
+from .windrose import DEFAULT_THETA_LABELS  # noqa
 from .windrose import DPI_DEFAULT  # noqa
 from .windrose import FIGSIZE_DEFAULT  # noqa
 from .windrose import WindAxes  # noqa

windrose/windrose.py

@@ -17,6 +17,7 @@
 DPI_DEFAULT = 80
 CALM_CIRCLE_COLOR = "red"
 CALM_CIRCLE_ALPHA = 0.4
+DEFAULT_THETA_LABELS = ["E", "N-E", "N", "N-W", "W", "S-W", "S", "S-E"]
 
 
 class WindAxesFactory:
@@ -70,24 +71,26 @@ class WindroseAxes(PolarAxes):
     def __init__(self, *args, **kwargs):
         """
         See Axes base class for args and kwargs documentation
+
+        Other kwargs are:
+
+        theta_labels : default ["E", "N-E", "N", "N-W", "W", "S-W", "S", "S-E"]
+            Labels for theta coordinate
         """
 
         # Uncomment to have the possibility to change the resolution directly
         # when the instance is created
         # self.RESOLUTION = kwargs.pop('resolution', 100)
         self.rmax = kwargs.pop("rmax", None)
-        self.theta_labels = kwargs.pop("theta_labels", None)
-        if self.theta_labels is None:
-            self.theta_labels = ["E", "N-E", "N", "N-W", "W", "S-W", "S", "S-E"]
+        self.theta_labels = kwargs.pop("theta_labels", DEFAULT_THETA_LABELS)
+
         PolarAxes.__init__(self, *args, **kwargs)
         self.set_aspect("equal", adjustable="box", anchor="C")
         self.radii_angle = 67.5
         self.clear()
 
     @staticmethod
-    def from_ax(
-        ax=None, fig=None, rmax=None, theta_labels=None, rect=None, *args, **kwargs
-    ):
+    def from_ax(ax=None, fig=None, rmax=None, rect=None, *args, **kwargs):
         """
         Return a WindroseAxes object for the figure `fig`.
         """
@@ -101,9 +104,7 @@ def from_ax(
                 )
             if rect is None:
                 rect = [0.1, 0.1, 0.8, 0.8]
-            ax = WindroseAxes(
-                fig, rect, rmax=rmax, theta_labels=theta_labels, *args, **kwargs
-            )
+            ax = WindroseAxes(fig, rect, rmax=rmax, *args, **kwargs)
             fig.add_axes(ax)
             return ax
         else: