CLN/ENH: clean up docstrings remove some cruft and generalize

Author: cpcloud

pandas/core/reshape.py

@@ -5,19 +5,20 @@
 
 import numpy as np
 
+import six
+
 from pandas.core.series import Series
 from pandas.core.frame import DataFrame
 
 from pandas.core.categorical import Categorical
 from pandas.core.common import (notnull, _ensure_platform_int, _maybe_promote,
-                                _maybe_upcast, isnull)
+                                isnull)
 from pandas.core.groupby import (get_group_index, _compress_group_index,
                                  decons_group_index)
 import pandas.core.common as com
 import pandas.algos as algos
-from pandas import lib
 
-from pandas.core.index import MultiIndex, Index
+from pandas.core.index import MultiIndex
 
 
 class ReshapeError(Exception):
@@ -35,21 +36,26 @@ class _Unstacker(object):
 
     Examples
     --------
+    >>> import pandas as pd
+    >>> index = pd.MultiIndex.from_tuples([('one', 'a'), ('one', 'b'),
+    ...                                    ('two', 'a'), ('two', 'b')])
+    >>> s = pd.Series(np.arange(1.0, 5.0), index=index)
     >>> s
-    one  a   1.
-    one  b   2.
-    two  a   3.
-    two  b   4.
+    one  a   1
+         b   2
+    two  a   3
+         b   4
+    dtype: float64
 
     >>> s.unstack(level=-1)
          a   b
-    one  1.  2.
-    two  3.  4.
+    one  1  2
+    two  3  4
 
     >>> s.unstack(level=0)
        one  two
-    a  1.   2.
-    b  3.   4.
+    a  1   2
+    b  3   4
 
     Returns
     -------
@@ -159,7 +165,7 @@ def get_result(self):
                     values[j] = orig_values[i]
             else:
                 index = index.take(self.unique_groups)
-        
+
         return DataFrame(values, index=index, columns=columns)
 
     def get_new_values(self):
@@ -617,9 +623,10 @@ def melt(frame, id_vars=None, value_vars=None,
 
     Examples
     --------
+    >>> import pandas as pd
     >>> df = pd.DataFrame({'A': {0: 'a', 1: 'b', 2: 'c'},
-              'B': {0: 1, 1: 3, 2: 5},
-              'C': {0: 2, 1: 4, 2: 6}})
+    ...                    'B': {0: 1, 1: 3, 2: 5},
+    ...                    'C': {0: 2, 1: 4, 2: 6}})
 
     >>> df
        A  B  C
@@ -632,7 +639,7 @@ def melt(frame, id_vars=None, value_vars=None,
     0  a        B      1
     1  b        B      3
     2  c        B      5
-    
+
     >>> melt(df, id_vars=['A'], value_vars=['B'],
     ... var_name='myVarname', value_name='myValname')
        A myVarname  myValname
@@ -679,9 +686,13 @@ def melt(frame, id_vars=None, value_vars=None,
             if len(frame.columns.names) == len(set(frame.columns.names)):
                 var_name = frame.columns.names
             else:
-                var_name = ['variable_%s' % i for i in range(len(frame.columns.names))]
+                var_name = ['variable_%s' % i for i in
+                            xrange(len(frame.columns.names))]
         else:
-            var_name = frame.columns.name if frame.columns.name is not None else 'variable'
+            var_name = [frame.columns.name if frame.columns.name is not None
+                        else 'variable']
+    if isinstance(var_name, six.string_types):
+        var_name = [var_name]
 
     N, K = frame.shape
     K -= len(id_vars)
@@ -690,17 +701,12 @@ def melt(frame, id_vars=None, value_vars=None,
     for col in id_vars:
         mdata[col] = np.tile(frame.pop(col).values, K)
 
-    if isinstance(var_name, list):
-        mcolumns = id_vars + var_name + [value_name]
-    else:
-        mcolumns = id_vars + [var_name, value_name]
+    mcolumns = id_vars + var_name + [value_name]
 
     mdata[value_name] = frame.values.ravel('F')
-    if isinstance(frame.columns, MultiIndex):
-        for i, col in enumerate(var_name):
-            mdata[col] = np.asarray(frame.columns.get_level_values(i)).repeat(N)
-    else: # assume isinstance(frame.columns, Index):
-        mdata[var_name] = np.asarray(frame.columns).repeat(N)
+    for i, col in enumerate(var_name):
+        # asanyarray will keep the columns as an Index
+        mdata[col] = np.asanyarray(frame.columns.get_level_values(i)).repeat(N)
 
     return DataFrame(mdata, columns=mcolumns)
 
@@ -718,13 +724,16 @@ def lreshape(data, groups, dropna=True, label=None):
 
     Examples
     --------
+    >>> import pandas as pd
+    >>> data = pd.DataFrame({'hr1': [514, 573], 'hr2': [545, 526],
+    ...                      'team': ['Red Sox', 'Yankees'],
+    ...                      'year1': [2007, 2008], 'year2': [2008, 2008]})
     >>> data
        hr1  hr2     team  year1  year2
     0  514  545  Red Sox   2007   2008
     1  573  526  Yankees   2007   2008
 
-    >>> pd.lreshape(data, {'year': ['year1', 'year2'],
-                           'hr': ['hr1', 'hr2']})
+    >>> pd.lreshape(data, {'year': ['year1', 'year2'], 'hr': ['hr1', 'hr2']})
           team   hr  year
     0  Red Sox  514  2007
     1  Yankees  573  2007