Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pandas/io/_util.py: 16%

1from __future__ import annotations 

2 

3import contextlib 

4import datetime as dt 

5from typing import ( 

6 TYPE_CHECKING, 

7 Any, 

8 Literal, 

9 cast, 

10) 

11import warnings 

12import zoneinfo 

13 

14import numpy as np 

15 

16from pandas._config import using_string_dtype 

17 

18from pandas._libs import lib 

19from pandas._libs.tslibs import timezones 

20from pandas.compat import ( 

21 pa_version_under18p0, 

22 pa_version_under19p0, 

23 pa_version_under25p0, 

24) 

25from pandas.compat._optional import import_optional_dependency 

26from pandas.errors import Pandas4Warning 

27 

28from pandas.core.dtypes.common import pandas_dtype 

29 

30import pandas as pd 

31 

32if TYPE_CHECKING: 

33 from collections.abc import ( 

34 Callable, 

35 Generator, 

36 Hashable, 

37 Sequence, 

38 ) 

39 

40 import pyarrow 

41 

42 from pandas._typing import ( 

43 DtypeArg, 

44 DtypeBackend, 

45 ) 

46 

47 from pandas.core.dtypes.base import ExtensionDtype 

48 

49 

50pytz = import_optional_dependency("pytz", errors="ignore") 

51 

52 

53@contextlib.contextmanager 

54def suppress_pyarrow_values_warning() -> Generator[None]: 

55 """ 

56 Suppress the deprecation warning pyarrow triggers by calling ``.values`` 

57 on timezone-aware data when converting pandas objects. 

58 

59 Remove once the minimum pyarrow (26.0) no longer does this, see GH#68426 and 

60 apache/arrow#51302. 

61 """ 

62 with warnings.catch_warnings(): 

63 warnings.filterwarnings( 

64 "ignore", 

65 r"(Series|DatetimeIndex)\.values returning an ndarray that drops " 

66 "timezone information", 

67 Pandas4Warning, 

68 ) 

69 yield 

70 

71 

72def _arrow_dtype_mapping() -> dict[pyarrow.DataType, Any]: 

73 pa = import_optional_dependency("pyarrow") 

74 return { 

75 pa.int8(): pd.Int8Dtype(), 

76 pa.int16(): pd.Int16Dtype(), 

77 pa.int32(): pd.Int32Dtype(), 

78 pa.int64(): pd.Int64Dtype(), 

79 pa.uint8(): pd.UInt8Dtype(), 

80 pa.uint16(): pd.UInt16Dtype(), 

81 pa.uint32(): pd.UInt32Dtype(), 

82 pa.uint64(): pd.UInt64Dtype(), 

83 pa.bool_(): pd.BooleanDtype(), 

84 pa.string(): pd.StringDtype(), 

85 pa.float32(): pd.Float32Dtype(), 

86 pa.float64(): pd.Float64Dtype(), 

87 pa.string(): pd.StringDtype(), 

88 pa.large_string(): pd.StringDtype(), 

89 } 

90 

91 

92def _arrow_string_types_mapper() -> Callable[[pyarrow.DataType], ExtensionDtype | None]: 

93 pa = import_optional_dependency("pyarrow") 

94 

95 mapping = { 

96 pa.string(): pd.StringDtype(na_value=np.nan), 

97 pa.large_string(): pd.StringDtype(na_value=np.nan), 

98 } 

99 if not pa_version_under18p0: 

100 mapping[pa.string_view()] = pd.StringDtype(na_value=np.nan) 

101 

102 return mapping.get 

103 

104 

105def arrow_table_to_pandas( 

106 table: pyarrow.Table, 

107 dtype_backend: DtypeBackend | Literal["numpy"] | lib.NoDefault = lib.no_default, 

108 null_to_int64: bool = False, 

109 to_pandas_kwargs: dict[str, Any] | None = None, 

110 dtype: DtypeArg | None = None, 

111 names: Sequence[Hashable] | None = None, 

112) -> pd.DataFrame: 

113 pa = import_optional_dependency("pyarrow") 

114 

115 to_pandas_kwargs = {} if to_pandas_kwargs is None else to_pandas_kwargs 

116 

117 types_mapper: ( 

118 type[pd.ArrowDtype] | Callable[[pyarrow.DataType], ExtensionDtype | None] | None 

119 ) 

120 if dtype_backend == "numpy_nullable": 

121 mapping = _arrow_dtype_mapping() 

122 if null_to_int64: 

123 # Modify the default mapping to also map null to Int64 

124 # (to match other engines - only for CSV parser) 

125 mapping[pa.null()] = pd.Int64Dtype() 

126 types_mapper = mapping.get 

127 elif dtype_backend == "pyarrow": 

128 types_mapper = pd.ArrowDtype 

129 elif using_string_dtype(): 

130 if pa_version_under19p0: 

131 types_mapper = _arrow_string_types_mapper() 

132 elif dtype is not None: 

133 # GH#56136 Avoid lossy conversion to float64 

134 # We'll convert to numpy below if 

135 types_mapper = { 

136 pa.int8(): pd.Int8Dtype(), 

137 pa.int16(): pd.Int16Dtype(), 

138 pa.int32(): pd.Int32Dtype(), 

139 pa.int64(): pd.Int64Dtype(), 

140 }.get 

141 else: 

142 types_mapper = None 

143 elif dtype_backend is lib.no_default or dtype_backend == "numpy": 

144 if dtype is not None: 

145 # GH#56136 Avoid lossy conversion to float64 

146 # We'll convert to numpy below if 

147 types_mapper = { 

148 pa.int8(): pd.Int8Dtype(), 

149 pa.int16(): pd.Int16Dtype(), 

150 pa.int32(): pd.Int32Dtype(), 

151 pa.int64(): pd.Int64Dtype(), 

152 }.get 

153 else: 

154 types_mapper = None 

155 else: 

156 raise NotImplementedError 

157 

158 df = table.to_pandas(types_mapper=types_mapper, **to_pandas_kwargs) 

159 df = _post_convert_dtypes(df, dtype_backend, dtype, names) 

160 df = _normalize_timezone_dtypes(df) 

161 return df 

162 

163 

164def _post_convert_dtypes( 

165 df: pd.DataFrame, 

166 dtype_backend: DtypeBackend | Literal["numpy"] | lib.NoDefault, 

167 dtype: DtypeArg | None, 

168 names: Sequence[Hashable] | None, 

169) -> pd.DataFrame: 

170 if dtype is not None and ( 

171 dtype_backend is lib.no_default or dtype_backend == "numpy" 

172 ): 

173 # GH#56136 apply any user-provided dtype, and convert any IntegerDtype 

174 # columns the user didn't explicitly ask for. 

175 if isinstance(dtype, dict): 

176 if names is not None: 

177 df.columns = names 

178 

179 cmp_dtypes = { 

180 pd.Int8Dtype(), 

181 pd.Int16Dtype(), 

182 pd.Int32Dtype(), 

183 pd.Int64Dtype(), 

184 } 

185 for col in df.columns: 

186 if col not in dtype and df[col].dtype in cmp_dtypes: 

187 # Any key that the user didn't explicitly specify 

188 # that got converted to IntegerDtype now gets converted 

189 # to numpy dtype. 

190 dtype[col] = df[col].dtype.numpy_dtype 

191 

192 # Ignore non-existent columns from dtype mapping 

193 # like other parsers do 

194 dtype = { 

195 key: pandas_dtype(dtype[key]) for key in dtype if key in df.columns 

196 } 

197 

198 # GH#65056 pyarrow returns datetime.date objects which trigger 

199 # a deprecation warning when compared against DatetimeIndex 

200 # categories. Convert to datetime64 first since the user can't 

201 # control pyarrow's behavior. 

202 for col, col_dtype in dtype.items(): 

203 if ( 

204 isinstance(col_dtype, pd.CategoricalDtype) 

205 and col_dtype.categories is not None 

206 and isinstance(col_dtype.categories, pd.DatetimeIndex) 

207 and df[col].dtype == np.dtype("object") 

208 and lib.infer_dtype(df[col]) == "date" 

209 ): 

210 df[col] = pd.to_datetime(df[col]) 

211 

212 else: 

213 dtype = pandas_dtype(dtype) 

214 

215 try: 

216 df = df.astype(dtype) 

217 except TypeError as err: 

218 # GH#44901 reraise to keep api consistent 

219 raise ValueError(str(err)) from err 

220 

221 # GH#56136 IntegerDtype was used above to avoid lossy float64 

222 # conversion in pyarrow; convert back to numpy now that the data 

223 # is categorical 

224 # runtime import to avoid circular import; core.dtypes.cast imports 

225 # this module at module scope 

226 from pandas.core.arrays.integer import IntegerDtype 

227 

228 col_dtypes = df.dtypes 

229 for i in range(len(df.columns)): 

230 col_dtype = col_dtypes.iloc[i] 

231 if not isinstance(col_dtype, pd.CategoricalDtype): 

232 continue 

233 cat_arr_dtype = col_dtype.categories.dtype 

234 if not isinstance(cat_arr_dtype, IntegerDtype): 

235 continue 

236 if isinstance(dtype, dict): 

237 requested = dtype.get(df.columns[i]) 

238 else: 

239 requested = dtype 

240 if ( 

241 isinstance(requested, pd.CategoricalDtype) 

242 and requested.categories is not None 

243 ): 

244 # the user explicitly asked for these categories 

245 continue 

246 new_cat_dtype = pd.CategoricalDtype( 

247 categories=col_dtype.categories.astype(cat_arr_dtype.numpy_dtype), 

248 ordered=col_dtype.ordered, 

249 ) 

250 df.isetitem(i, df.iloc[:, i].astype(new_cat_dtype)) 

251 

252 if ( 

253 not using_string_dtype() 

254 and dtype != "str" 

255 and (dtype_backend is lib.no_default or dtype_backend == "numpy") 

256 ): 

257 # Convert any StringDtype columns back to object dtype (pyarrow always 

258 # uses string dtype even when the infer_string option is False) 

259 for i in range(len(df.columns)): 

260 new_col = _maybe_convert_string_to_object(df.iloc[:, i]) 

261 if new_col is not None: 

262 df.isetitem(i, new_col) 

263 

264 new_idx = _maybe_convert_string_index_to_object(df.index) 

265 if new_idx is not None: 

266 df.index = new_idx 

267 new_cols = _maybe_convert_string_index_to_object(df.columns) 

268 if new_cols is not None: 

269 df.columns = new_cols 

270 

271 return df 

272 

273 

274def _maybe_convert_string_to_object( 

275 data: pd.Series | pd.Index, 

276) -> pd.Series | pd.Index | None: 

277 if isinstance(data.dtype, pd.StringDtype) and data.dtype.na_value is np.nan: 

278 return data.astype("object").fillna(None) 

279 elif isinstance(data.dtype, pd.CategoricalDtype): 

280 cat_dtype = data.dtype.categories.dtype 

281 if isinstance(cat_dtype, pd.StringDtype) and cat_dtype.na_value is np.nan: 

282 # not astype: for ordered categoricals CategoricalDtype.__eq__ 

283 # ignores the categories' dtype, so astype would no-op 

284 cat_arr = cast("pd.Categorical", data._values) 

285 new_arr = cat_arr.set_categories(data.dtype.categories.astype("object")) 

286 if isinstance(data, pd.Index): 

287 return pd.CategoricalIndex(new_arr, name=data.name) 

288 return pd.Series(new_arr, index=data.index, name=data.name, copy=False) 

289 

290 # no conversion needed 

291 return None 

292 

293 

294def _maybe_convert_string_index_to_object(index: pd.Index) -> pd.Index | None: 

295 if isinstance(index, pd.MultiIndex): 

296 if any( 

297 isinstance(level.dtype, pd.StringDtype) and level.dtype.na_value is np.nan 

298 for level in index.levels 

299 ): 

300 new_levels = [] 

301 for level in index.levels: 

302 new_level = _maybe_convert_string_to_object(level) 

303 if new_level is not None: 

304 new_levels.append(new_level) 

305 else: 

306 new_levels.append(level) 

307 return index.set_levels(new_levels) 

308 return None 

309 

310 else: 

311 return cast("pd.Index | None", _maybe_convert_string_to_object(index)) 

312 

313 

314def _normalize_pytz_timezone(tz: dt.tzinfo) -> dt.tzinfo: 

315 """ 

316 If the input tz is a pytz timezone, attempt to convert it to "default" 

317 tzinfo object (zoneinfo or datetime.timezone). 

318 """ 

319 if not type(tz).__module__.startswith("pytz"): 

320 # isinstance(col.dtype.tz, pytz.BaseTzInfo) does not included 

321 # fixed offsets 

322 return tz 

323 

324 if timezones.is_utc(tz): 

325 return dt.UTC 

326 

327 if tz.zone is not None: # type: ignore[attr-defined] 

328 try: 

329 return zoneinfo.ZoneInfo(tz.zone) # type: ignore[attr-defined] 

330 except Exception: 

331 # some pytz timezones might not be available for zoneinfo 

332 pass 

333 

334 if timezones.is_fixed_offset(tz): 

335 # Convert pytz fixed offset to datetime.timezone 

336 try: 

337 offset = tz.utcoffset(None) 

338 if offset is not None: 

339 return dt.timezone(offset) 

340 except Exception: 

341 pass 

342 

343 return tz 

344 

345 

346def _normalize_timezone_index(index: pd.Index) -> pd.Index: 

347 if isinstance(index, pd.MultiIndex): 

348 if any(isinstance(level.dtype, pd.DatetimeTZDtype) for level in index.levels): 

349 levels = [_normalize_timezone_index(level) for level in index.levels] 

350 return index.set_levels(levels) 

351 

352 return index 

353 

354 if isinstance(index.dtype, pd.DatetimeTZDtype): 

355 normalized_tz = _normalize_pytz_timezone(index.dtype.tz) 

356 if normalized_tz is not index.dtype.tz: 

357 return index.tz_convert(normalized_tz) # type: ignore[attr-defined] 

358 

359 return index 

360 

361 

362def _normalize_timezone_dtypes(df: pd.DataFrame) -> pd.DataFrame: 

363 """ 

364 PyArrow below version 25 uses pytz by default for timezones, but pandas uses 

365 zoneinfo / datetime.timezone since pandas 3.0. 

366 

367 Can be dropped once the minimum supported pyarrow is 25, which uses zoneinfo 

368 itself (https://github.com/apache/arrow/pull/49694). 

369 """ 

370 if pytz is not None and pa_version_under25p0: 

371 # Convert any pytz timezones to zoneinfo / fixed offset timezones 

372 if any( 

373 isinstance(dtype, pd.DatetimeTZDtype) 

374 for dtype in df._mgr.get_unique_dtypes() 

375 ): 

376 col_indices = df._select_dtypes_indices(pd.DatetimeTZDtype) 

377 for i in col_indices: 

378 col = df.iloc[:, i] 

379 normalized_tz = _normalize_pytz_timezone(col.dtype.tz) 

380 if normalized_tz is not col.dtype.tz: 

381 df.isetitem(i, col.dt.tz_convert(normalized_tz)) 

382 

383 df.index = _normalize_timezone_index(df.index) 

384 df.columns = _normalize_timezone_index(df.columns) 

385 

386 return df