ENH: tweak na_option checking per #2159

Author: wesm

RELEASE.rst

@@ -30,6 +30,7 @@ pandas 0.9.1
 **New features**
 
   - Can specify multiple sort orders in DataFrame/Series.sort/sort_index (#928)
+  - New `top` and `bottom` options for handling NAs in rank (#1508, #2159)
 
 **Improvements to existing features**
 
@@ -50,6 +51,7 @@ pandas 0.9.1
   - Fix Series.str failures when using pipe pattern '|' (#2119)
   - Fix pretty-printing of dict entries in Series, DataFrame (#2144)
   - Cast other datetime64 values to nanoseconds in DataFrame ctor (#2095)
+  - Alias Timestamp.astimezone to tz_convert, so will yield Timestamp (#2060)
 
 pandas 0.9.0
 ============

pandas/src/stats.pyx

@@ -83,10 +83,13 @@ def rank_1d_float64(object in_arr, ties_method='average', ascending=True,
         float64_t val, nan_value
         float64_t sum_ranks = 0
         int tiebreak = 0
+        bint keep_na = 0
     tiebreak = tiebreakers[ties_method]
 
     values = np.asarray(in_arr).copy()
 
+    keep_na = na_option == 'keep'
+
     if ascending ^ (na_option == 'top'):
         nan_value = np.inf
     else:
@@ -116,7 +119,7 @@ def rank_1d_float64(object in_arr, ties_method='average', ascending=True,
         sum_ranks += i + 1
         dups += 1
         val = sorted_data[i]
-        if (val == nan_value) and (na_option == 'keep'):
+        if (val == nan_value) and keep_na:
             ranks[argsorted[i]] = nan
             continue
         if i == n - 1 or fabs(sorted_data[i + 1] - val) > FP_ERR:
@@ -212,8 +215,12 @@ def rank_2d_float64(object in_arr, axis=0, ties_method='average',
         float64_t val, nan_value
         float64_t sum_ranks = 0
         int tiebreak = 0
+        bint keep_na = 0
+
     tiebreak = tiebreakers[ties_method]
 
+    keep_na = na_option == 'keep'
+
     in_arr = np.asarray(in_arr)
 
     if axis == 0:
@@ -251,7 +258,7 @@ def rank_2d_float64(object in_arr, axis=0, ties_method='average',
             sum_ranks += j + 1
             dups += 1
             val = values[i, j]
-            if val == nan_value and na_option == 'keep':
+            if val == nan_value and keep_na:
                 ranks[i, argsorted[i, j]] = nan
                 continue
             if j == k - 1 or fabs(values[i, j + 1] - val) > FP_ERR:
@@ -360,8 +367,12 @@ def rank_1d_generic(object in_arr, bint retry=1, ties_method='average',
         object val, nan_value
         float64_t sum_ranks = 0
         int tiebreak = 0
+        bint keep_na = 0
+
     tiebreak = tiebreakers[ties_method]
 
+    keep_na = na_option == 'keep'
+
     values = np.array(in_arr, copy=True)
 
     if values.dtype != np.object_:
@@ -403,7 +414,7 @@ def rank_1d_generic(object in_arr, bint retry=1, ties_method='average',
         sum_ranks += i + 1
         dups += 1
         val = util.get_value_at(sorted_data, i)
-        if val is nan_value and na_option=='keep':
+        if val is nan_value and keep_na:
             ranks[argsorted[i]] = nan
             continue
         if (i == n - 1 or
@@ -465,8 +476,12 @@ def rank_2d_generic(object in_arr, axis=0, ties_method='average',
         object val, nan_value
         float64_t sum_ranks = 0
         int tiebreak = 0
+        bint keep_na = 0
+
     tiebreak = tiebreakers[ties_method]
 
+    keep_na = na_option == 'keep'
+
     in_arr = np.asarray(in_arr)
 
     if axis == 0:
@@ -512,7 +527,7 @@ def rank_2d_generic(object in_arr, axis=0, ties_method='average',
         dups = sum_ranks = infs = 0
         for j in range(k):
             val = values[i, j]
-            if val is nan_value and na_option == 'keep':
+            if val is nan_value and keep_na:
                 ranks[i, argsorted[i, j]] = nan
                 infs += 1
                 continue