@@ -801,6 +801,26 @@ def test_pivot_table_margins_name_with_aggfunc_list(self):
801801 expected = pd .DataFrame (table .values , index = ix , columns = cols )
802802 tm .assert_frame_equal (table , expected )
803803
804+ def test_categorical_margins (self ):
805+ # GH 10989
806+ df = pd .DataFrame ({'x' : np .arange (8 ),
807+ 'y' : np .arange (8 ) // 4 ,
808+ 'z' : np .arange (8 ) % 2 })
809+
810+ expected = pd .DataFrame ([[1.0 , 2.0 , 1.5 ], [5 , 6 , 5.5 ], [3 , 4 , 3.5 ]])
811+ expected .index = Index ([0 , 1 , 'All' ], name = 'y' )
812+ expected .columns = Index ([0 , 1 , 'All' ], name = 'z' )
813+
814+ data = df .copy ()
815+ table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
816+ tm .assert_frame_equal (table , expected )
817+
818+ data = df .copy ()
819+ data .y = data .y .astype ('category' )
820+ data .z = data .z .astype ('category' )
821+ table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
822+ tm .assert_frame_equal (table , expected )
823+
804824
805825class TestCrosstab (tm .TestCase ):
806826
@@ -919,26 +939,6 @@ def test_crosstab_dropna(self):
919939 names = ['b' , 'c' ])
920940 tm .assert_index_equal (res .columns , m )
921941
922- def test_categorical_margins (self ):
923- # GH 10989
924- df = pd .DataFrame ({'x' : np .arange (8 ),
925- 'y' : np .arange (8 ) // 4 ,
926- 'z' : np .arange (8 ) % 2 })
927-
928- expected = pd .DataFrame ([[1.0 , 2.0 , 1.5 ], [5 , 6 , 5.5 ], [3 , 4 , 3.5 ]])
929- expected .index = Index ([0 , 1 , 'All' ], name = 'y' )
930- expected .columns = Index ([0 , 1 , 'All' ], name = 'z' )
931-
932- data = df .copy ()
933- table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
934- tm .assert_frame_equal (table , expected )
935-
936- data = df .copy ()
937- data .y = data .y .astype ('category' )
938- data .z = data .z .astype ('category' )
939- table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
940- tm .assert_frame_equal (table , expected )
941-
942942 def test_crosstab_no_overlap (self ):
943943 # GS 10291
944944
0 commit comments