@@ -1941,18 +1941,19 @@ def test_chunked_categorical(version):
1941
1941
1942
1942
def test_chunked_categorical_partial (dirpath ):
1943
1943
dta_file = os .path .join (dirpath , "stata-dta-partially-labeled.dta" )
1944
- reader = StataReader (dta_file , chunksize = 2 )
1945
1944
values = ["a" , "b" , "a" , "b" , 3.0 ]
1946
- with pytest .warns (CategoricalConversionWarning , match = "One or more series" ):
1947
- for i , block in enumerate (reader ):
1948
- assert list (block .cats ) == values [2 * i : 2 * (i + 1 )]
1949
- if i < 2 :
1950
- idx = pd .Index (["a" , "b" ])
1951
- else :
1952
- idx = pd .Float64Index ([3.0 ])
1953
- tm .assert_index_equal (block .cats .cat .categories , idx )
1954
- reader = StataReader (dta_file , chunksize = 5 )
1955
- large_chunk = reader .__next__ ()
1945
+ with StataReader (dta_file , chunksize = 2 ) as reader :
1946
+ with tm .assert_produces_warning (CategoricalConversionWarning ):
1947
+ for i , block in enumerate (reader ):
1948
+ assert list (block .cats ) == values [2 * i : 2 * (i + 1 )]
1949
+ if i < 2 :
1950
+ idx = pd .Index (["a" , "b" ])
1951
+ else :
1952
+ idx = pd .Float64Index ([3.0 ])
1953
+ tm .assert_index_equal (block .cats .cat .categories , idx )
1954
+ with tm .assert_produces_warning (CategoricalConversionWarning ):
1955
+ with StataReader (dta_file , chunksize = 5 ) as reader :
1956
+ large_chunk = reader .__next__ ()
1956
1957
direct = read_stata (dta_file )
1957
1958
tm .assert_frame_equal (direct , large_chunk )
1958
1959
@@ -1966,4 +1967,5 @@ def test_iterator_errors(dirpath):
1966
1967
with pytest .raises (ValueError , match = "chunksize must be a positive" ):
1967
1968
StataReader (dta_file , chunksize = "apple" )
1968
1969
with pytest .raises (ValueError , match = "chunksize must be set to a positive" ):
1969
- StataReader (dta_file ).__next__ ()
1970
+ with StataReader (dta_file ) as reader :
1971
+ reader .__next__ ()
0 commit comments