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

1from __future__ import annotations 

2 

3import datetime as dt 

4from typing import ( 

5 TYPE_CHECKING, 

6 Literal, 

7 cast, 

8) 

9import zoneinfo 

10 

11import numpy as np 

12 

13from pandas._config import using_string_dtype 

14 

15from pandas._libs import lib 

16from pandas._libs.tslibs import timezones 

17from pandas.compat import ( 

18 pa_version_under18p0, 

19 pa_version_under19p0, 

20) 

21from pandas.compat._optional import import_optional_dependency 

22 

23from pandas.core.dtypes.common import pandas_dtype 

24 

25import pandas as pd 

26 

27if TYPE_CHECKING: 

28 from collections.abc import ( 

29 Callable, 

30 Hashable, 

31 Sequence, 

32 ) 

33 

34 import pyarrow 

35 

36 from pandas._typing import ( 

37 DtypeArg, 

38 DtypeBackend, 

39 ) 

40 

41 

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

43 

44 

45def _arrow_dtype_mapping() -> dict: 

46 pa = import_optional_dependency("pyarrow") 

47 return { 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

62 } 

63 

64 

65def _arrow_string_types_mapper() -> Callable: 

66 pa = import_optional_dependency("pyarrow") 

67 

68 mapping = { 

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

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

71 } 

72 if not pa_version_under18p0: 

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

74 

75 return mapping.get 

76 

77 

78def arrow_table_to_pandas( 

79 table: pyarrow.Table, 

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

81 null_to_int64: bool = False, 

82 to_pandas_kwargs: dict | None = None, 

83 dtype: DtypeArg | None = None, 

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

85) -> pd.DataFrame: 

86 pa = import_optional_dependency("pyarrow") 

87 

88 to_pandas_kwargs = {} if to_pandas_kwargs is None else to_pandas_kwargs 

89 

90 types_mapper: type[pd.ArrowDtype] | None | Callable 

91 if dtype_backend == "numpy_nullable": 

92 mapping = _arrow_dtype_mapping() 

93 if null_to_int64: 

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

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

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

97 types_mapper = mapping.get 

98 elif dtype_backend == "pyarrow": 

99 types_mapper = pd.ArrowDtype 

100 elif using_string_dtype(): 

101 if pa_version_under19p0: 

102 types_mapper = _arrow_string_types_mapper() 

103 elif dtype is not None: 

104 # GH#56136 Avoid lossy conversion to float64 

105 # We'll convert to numpy below if 

106 types_mapper = { 

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

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

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

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

111 }.get 

112 else: 

113 types_mapper = None 

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

115 if dtype is not None: 

116 # GH#56136 Avoid lossy conversion to float64 

117 # We'll convert to numpy below if 

118 types_mapper = { 

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

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

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

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

123 }.get 

124 else: 

125 types_mapper = None 

126 else: 

127 raise NotImplementedError 

128 

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

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

131 df = _normalize_timezone_dtypes(df) 

132 return df 

133 

134 

135def _post_convert_dtypes( 

136 df: pd.DataFrame, 

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

138 dtype: DtypeArg | None, 

139 names: Sequence[Hashable] | None, 

140) -> pd.DataFrame: 

141 if dtype is not None and ( 

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

143 ): 

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

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

146 if isinstance(dtype, dict): 

147 if names is not None: 

148 df.columns = names 

149 

150 cmp_dtypes = { 

151 pd.Int8Dtype(), 

152 pd.Int16Dtype(), 

153 pd.Int32Dtype(), 

154 pd.Int64Dtype(), 

155 } 

156 for col in df.columns: 

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

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

159 # that got converted to IntegerDtype now gets converted 

160 # to numpy dtype. 

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

162 

163 # Ignore non-existent columns from dtype mapping 

164 # like other parsers do 

165 dtype = { 

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

167 } 

168 

169 else: 

170 dtype = pandas_dtype(dtype) 

171 

172 try: 

173 df = df.astype(dtype) 

174 except TypeError as err: 

175 # GH#44901 reraise to keep api consistent 

176 raise ValueError(str(err)) from err 

177 

178 if ( 

179 not using_string_dtype() 

180 and dtype != "str" 

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

182 ): 

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

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

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

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

187 if new_col is not None: 

188 df.isetitem(i, new_col) 

189 

190 new_idx = _maybe_convert_string_index_to_object(df.index) 

191 if new_idx is not None: 

192 df.index = new_idx 

193 new_cols = _maybe_convert_string_index_to_object(df.columns) 

194 if new_cols is not None: 

195 df.columns = new_cols 

196 

197 return df 

198 

199 

200def _maybe_convert_string_to_object( 

201 data: pd.Series | pd.Index, 

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

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

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

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

206 cat_dtype = data.dtype.categories.dtype 

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

208 cat_dtype = pd.CategoricalDtype( 

209 categories=data.dtype.categories.astype("object"), 

210 ordered=data.dtype.ordered, 

211 ) 

212 return data.astype(cat_dtype) 

213 

214 # no conversion needed 

215 return None 

216 

217 

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

219 if isinstance(index, pd.MultiIndex): 

220 if any( 

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

222 for level in index.levels 

223 ): 

224 new_levels = [] 

225 for level in index.levels: 

226 new_level = _maybe_convert_string_to_object(level) 

227 if new_level is not None: 

228 new_levels.append(new_level) 

229 else: 

230 new_levels.append(level) 

231 return index.set_levels(new_levels) 

232 return None 

233 

234 else: 

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

236 

237 

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

239 """ 

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

241 tzinfo object (zoneinfo or datetime.timezone). 

242 """ 

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

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

245 # fixed offsets 

246 return tz 

247 

248 if timezones.is_utc(tz): 

249 return dt.timezone.utc 

250 

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

252 try: 

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

254 except Exception: 

255 # some pytz timezones might not be available for zoneinfo 

256 pass 

257 

258 if timezones.is_fixed_offset(tz): 

259 # Convert pytz fixed offset to datetime.timezone 

260 try: 

261 offset = tz.utcoffset(None) 

262 if offset is not None: 

263 return dt.timezone(offset) 

264 except Exception: 

265 pass 

266 

267 return tz 

268 

269 

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

271 if isinstance(index, pd.MultiIndex): 

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

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

274 return index.set_levels(levels) 

275 

276 return index 

277 

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

279 normalized_tz = _normalize_pytz_timezone(index.dtype.tz) 

280 if normalized_tz is not index.dtype.tz: 

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

282 

283 return index 

284 

285 

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

287 """ 

288 PyArrow uses pytz by default for timezones, but pandas uses 

289 zoneinfo / datetime.timezone since pandas 3.0. 

290 

291 TODO: Starting with pyarrow 25, it will use zoneinfo by default, and then 

292 this normalization can be skipped (https://github.com/apache/arrow/pull/49694). 

293 """ 

294 if pytz is not None: 

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

296 if any( 

297 isinstance(dtype, pd.DatetimeTZDtype) 

298 for dtype in df._mgr.get_unique_dtypes() 

299 ): 

300 col_indices = df._select_dtypes_indices(pd.DatetimeTZDtype) 

301 for i in col_indices: 

302 col = df.iloc[:, i] 

303 normalized_tz = _normalize_pytz_timezone(col.dtype.tz) 

304 if normalized_tz is not col.dtype.tz: 

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

306 

307 df.index = _normalize_timezone_index(df.index) 

308 df.columns = _normalize_timezone_index(df.columns) 

309 

310 return df