@@ -635,6 +635,57 @@ def test_get(self):
635
635
636
636
pytest .raises (KeyError , store .get , 'b' )
637
637
638
+ @pytest .mark .parametrize ('where, expected' , [
639
+ ('/' , {
640
+ '' : ({'first_group' , 'second_group' }, set ()),
641
+ '/first_group' : (set (), {'df1' , 'df2' }),
642
+ '/second_group' : ({'third_group' }, {'df3' , 's1' }),
643
+ '/second_group/third_group' : (set (), {'df4' }),
644
+ }),
645
+ ('/second_group' , {
646
+ '/second_group' : ({'third_group' }, {'df3' , 's1' }),
647
+ '/second_group/third_group' : (set (), {'df4' }),
648
+ })
649
+ ])
650
+ def test_walk (self , where , expected ):
651
+ # GH10143
652
+ objs = {
653
+ 'df1' : pd .DataFrame ([1 , 2 , 3 ]),
654
+ 'df2' : pd .DataFrame ([4 , 5 , 6 ]),
655
+ 'df3' : pd .DataFrame ([6 , 7 , 8 ]),
656
+ 'df4' : pd .DataFrame ([9 , 10 , 11 ]),
657
+ 's1' : pd .Series ([10 , 9 , 8 ]),
658
+ # Next 3 items aren't pandas objects and should be ignored
659
+ 'a1' : np .array ([[1 , 2 , 3 ], [4 , 5 , 6 ]]),
660
+ 'tb1' : np .array ([(1 , 2 , 3 ), (4 , 5 , 6 )], dtype = 'i,i,i' ),
661
+ 'tb2' : np .array ([(7 , 8 , 9 ), (10 , 11 , 12 )], dtype = 'i,i,i' )
662
+ }
663
+
664
+ with ensure_clean_store ('walk_groups.hdf' , mode = 'w' ) as store :
665
+ store .put ('/first_group/df1' , objs ['df1' ])
666
+ store .put ('/first_group/df2' , objs ['df2' ])
667
+ store .put ('/second_group/df3' , objs ['df3' ])
668
+ store .put ('/second_group/s1' , objs ['s1' ])
669
+ store .put ('/second_group/third_group/df4' , objs ['df4' ])
670
+ # Create non-pandas objects
671
+ store ._handle .create_array ('/first_group' , 'a1' , objs ['a1' ])
672
+ store ._handle .create_table ('/first_group' , 'tb1' , obj = objs ['tb1' ])
673
+ store ._handle .create_table ('/second_group' , 'tb2' , obj = objs ['tb2' ])
674
+
675
+ assert len (list (store .walk (where = where ))) == len (expected )
676
+ for path , groups , leaves in store .walk (where = where ):
677
+ assert path in expected
678
+ expected_groups , expected_frames = expected [path ]
679
+ assert expected_groups == set (groups )
680
+ assert expected_frames == set (leaves )
681
+ for leaf in leaves :
682
+ frame_path = '/' .join ([path , leaf ])
683
+ obj = store .get (frame_path )
684
+ if 'df' in leaf :
685
+ tm .assert_frame_equal (obj , objs [leaf ])
686
+ else :
687
+ tm .assert_series_equal (obj , objs [leaf ])
688
+
638
689
def test_getattr (self ):
639
690
640
691
with ensure_clean_store (self .path ) as store :
0 commit comments