@@ -1941,18 +1941,19 @@ def test_chunked_categorical(version):
19411941
19421942def test_chunked_categorical_partial (dirpath ):
19431943 dta_file = os .path .join (dirpath , "stata-dta-partially-labeled.dta" )
1944- reader = StataReader (dta_file , chunksize = 2 )
19451944 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__ ()
19561957 direct = read_stata (dta_file )
19571958 tm .assert_frame_equal (direct , large_chunk )
19581959
@@ -1966,4 +1967,5 @@ def test_iterator_errors(dirpath):
19661967 with pytest .raises (ValueError , match = "chunksize must be a positive" ):
19671968 StataReader (dta_file , chunksize = "apple" )
19681969 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