use apply()-kwargs instead of partial, more tests, better examples

AnkurDedania · Jan 2, 2017 · a0558e7 · a0558e7
1 parent c8d3ac4
commit a0558e7
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 9 deletions.
diff --git a/pandas/core/frame.py b/pandas/core/frame.py
@@ -4989,7 +4989,7 @@ def nunique(self, axis=0, dropna=True):
 
         Examples
         --------
-        >>> df = DataFrame({'A': [1, 2, 3], 'B': [1, 1, 1]})
+        >>> df = pd.DataFrame({'A': [1, 2, 3], 'B': [1, 1, 1]})
         >>> df.nunique()
         A    3
         B    1
@@ -4999,8 +4999,7 @@ def nunique(self, axis=0, dropna=True):
         1    2
         2    2
         """
-        func = functools.partial(Series.nunique, dropna=dropna)
-        return self.apply(func, axis=axis)
+        return self.apply(Series.nunique, axis=axis, dropna=dropna)
 
     def idxmin(self, axis=0, skipna=True):
         """

diff --git a/pandas/core/groupby.py b/pandas/core/groupby.py
@@ -3917,9 +3917,10 @@ def nunique(self, dropna=True):
 
         Examples
         --------
-        >>> df = DataFrame({'id': ['spam', 'egg', 'egg', 'spam', 'ham', 'ham'],
-        ...                 'value1': [1, 5, 5, 2, 5, 5],
-        ...                 'value2': list('abbaxy')})
+        >>> df = pd.DataFrame({'id': ['spam', 'egg', 'egg', 'spam',
+        ...                           'ham', 'ham'],
+        ...                    'value1': [1, 5, 5, 2, 5, 5],
+        ...                    'value2': list('abbaxy')})
         >>> df
              id  value1 value2
         0  spam       1      a
@@ -3936,16 +3937,15 @@ def nunique(self, dropna=True):
         ham    1       1       2
         spam   1       2       1
 
+        # check for rows with the same id but conflicting values
         >>> df.groupby('id').filter(lambda g: (g.nunique() > 1).any())
              id  value1 value2
         0  spam       1      a
         3  spam       2      a
         4   ham       5      x
         5   ham       5      y
         """
-        from functools import partial
-        func = partial(Series.nunique, dropna=dropna)
-        return self.apply(lambda g: g.apply(func))
+        return self.apply(lambda g: g.apply(Series.nunique, dropna=dropna))
 
 
 from pandas.tools.plotting import boxplot_frame_groupby  # noqa

diff --git a/pandas/tests/frame/test_analytics.py b/pandas/tests/frame/test_analytics.py
@@ -416,6 +416,16 @@ def test_nunique(self):
         self._check_stat_op('nunique', f, has_skipna=False,
                             check_dtype=False, check_dates=True)
 
+        df = DataFrame({'A': [1, 1, 1],
+                        'B': [1, 2, 3],
+                        'C': [1, np.nan, 3]})
+        tm.assert_series_equal(df.nunique(), Series({'A': 1, 'B': 3, 'C': 2}))
+        tm.assert_series_equal(df.nunique(dropna=False),
+                               Series({'A': 1, 'B': 3, 'C': 3}))
+        tm.assert_series_equal(df.nunique(axis=1), Series({0: 1, 1: 2, 2: 2}))
+        tm.assert_series_equal(df.nunique(axis=1, dropna=False),
+                               Series({0: 1, 1: 3, 2: 2}))
+
     def test_sum(self):
         self._check_stat_op('sum', np.sum, has_numeric_only=True)