1from __future__ import annotations
2
3from collections.abc import (
4 Callable,
5 Hashable,
6 Iterable,
7 MutableMapping,
8 Sequence,
9)
10from typing import (
11 TYPE_CHECKING,
12 Any,
13 Literal,
14 TypeGuard,
15 TypeVar,
16 overload,
17)
18
19from pandas.compat._optional import import_optional_dependency
20
21from pandas.core.dtypes.common import (
22 is_integer,
23 is_list_like,
24)
25
26if TYPE_CHECKING:
27 from pandas.io.excel._base import ExcelWriter
28
29 ExcelWriter_t = type[ExcelWriter[Any]]
30 usecols_func = TypeVar("usecols_func", bound=Callable[[Hashable], object])
31
32_writers: MutableMapping[str, ExcelWriter_t] = {}
33
34
35def register_writer(klass: ExcelWriter_t) -> None:
36 """
37 Add engine to the excel writer registry.io.excel.
38
39 You must use this method to integrate with ``to_excel``.
40
41 Parameters
42 ----------
43 klass : ExcelWriter
44 """
45 if not callable(klass):
46 raise ValueError("Can only register callables as engines")
47 engine_name = klass._engine
48 _writers[engine_name] = klass
49
50
51def get_default_engine(ext: str, mode: Literal["reader", "writer"] = "reader") -> str:
52 """
53 Return the default reader/writer for the given extension.
54
55 Parameters
56 ----------
57 ext : str
58 The excel file extension for which to get the default engine.
59 mode : str {'reader', 'writer'}
60 Whether to get the default engine for reading or writing.
61 Either 'reader' or 'writer'
62
63 Returns
64 -------
65 str
66 The default engine for the extension.
67 """
68 _default_readers = {
69 "xlsx": "openpyxl",
70 "xlsm": "openpyxl",
71 "xlsb": "pyxlsb",
72 "xls": "xlrd",
73 "ods": "odf",
74 }
75 _default_writers = {
76 "xlsx": "openpyxl",
77 "xlsm": "openpyxl",
78 "xlsb": "pyxlsb",
79 "ods": "odf",
80 }
81 assert mode in ["reader", "writer"]
82 if mode == "writer":
83 # Prefer xlsxwriter over openpyxl if installed
84 xlsxwriter = import_optional_dependency("xlsxwriter", errors="warn")
85 if xlsxwriter:
86 _default_writers["xlsx"] = "xlsxwriter"
87 return _default_writers[ext]
88 else:
89 # Fall back to calamine if the deprecated default engine
90 # is not installed (GH#56542)
91 if ext == "xls" and not import_optional_dependency("xlrd", errors="ignore"):
92 _default_readers["xls"] = "calamine"
93 if ext == "xlsb" and not import_optional_dependency("pyxlsb", errors="ignore"):
94 _default_readers["xlsb"] = "calamine"
95 return _default_readers[ext]
96
97
98def get_writer(engine_name: str) -> ExcelWriter_t:
99 try:
100 return _writers[engine_name]
101 except KeyError as err:
102 raise ValueError(f"No Excel writer '{engine_name}'") from err
103
104
105def _excel2num(x: str) -> int:
106 """
107 Convert Excel column name like 'AB' to 0-based column index.
108
109 Parameters
110 ----------
111 x : str
112 The Excel column name to convert to a 0-based column index.
113
114 Returns
115 -------
116 num : int
117 The column index corresponding to the name.
118
119 Raises
120 ------
121 ValueError
122 Part of the Excel column name was invalid.
123 """
124 index = 0
125
126 for c in x.upper().strip():
127 cp = ord(c)
128
129 if cp < ord("A") or cp > ord("Z"):
130 raise ValueError(f"Invalid column name: {x}")
131
132 index = index * 26 + cp - ord("A") + 1
133
134 return index - 1
135
136
137def _range2cols(areas: str) -> list[int]:
138 """
139 Convert comma separated list of column names and ranges to indices.
140
141 Parameters
142 ----------
143 areas : str
144 A string containing a sequence of column ranges (or areas).
145
146 Returns
147 -------
148 cols : list
149 A list of 0-based column indices.
150
151 Examples
152 --------
153 >>> _range2cols("A:E")
154 [0, 1, 2, 3, 4]
155 >>> _range2cols("A,C,Z:AB")
156 [0, 2, 25, 26, 27]
157 """
158 cols: list[int] = []
159
160 for rng in areas.split(","):
161 if ":" in rng:
162 rngs = rng.split(":")
163 cols.extend(range(_excel2num(rngs[0]), _excel2num(rngs[1]) + 1))
164 else:
165 cols.append(_excel2num(rng))
166
167 return cols
168
169
170@overload
171def maybe_convert_usecols(usecols: str | list[int]) -> list[int]: ...
172
173
174@overload
175def maybe_convert_usecols(usecols: list[str]) -> list[str]: ...
176
177
178@overload
179def maybe_convert_usecols(usecols: usecols_func) -> usecols_func: ...
180
181
182@overload
183def maybe_convert_usecols(usecols: None) -> None: ...
184
185
186def maybe_convert_usecols(
187 usecols: str | list[int] | list[str] | usecols_func | None,
188) -> list[int] | list[str] | usecols_func | None:
189 """
190 Convert `usecols` into a compatible format for parsing in `parsers.py`.
191
192 Parameters
193 ----------
194 usecols : object
195 The use-columns object to potentially convert.
196
197 Returns
198 -------
199 converted : object
200 The compatible format of `usecols`.
201 """
202 if usecols is None:
203 return usecols
204
205 if is_integer(usecols):
206 raise ValueError(
207 "Passing an integer for `usecols` is no longer supported. "
208 "Please pass in a list of int from 0 to `usecols` inclusive instead."
209 )
210
211 if isinstance(usecols, str):
212 return _range2cols(usecols)
213
214 return usecols
215
216
217def validate_freeze_panes(
218 freeze_panes: tuple[int, int] | None,
219) -> TypeGuard[tuple[int, int]]:
220 if freeze_panes is not None:
221 if len(freeze_panes) == 2 and all(
222 isinstance(item, int) for item in freeze_panes
223 ):
224 return True
225
226 raise ValueError(
227 "freeze_panes must be of form (row, column) "
228 "where row and column are integers"
229 )
230
231 # freeze_panes wasn't specified, return False so it won't be applied
232 # to output sheet
233 return False
234
235
236def fill_mi_header(
237 row: list[Hashable], control_row: list[bool]
238) -> tuple[list[Hashable], list[bool]]:
239 """
240 Forward fill blank entries in row but only inside the same parent index.
241
242 Used for creating headers in MultiIndex.
243
244 Parameters
245 ----------
246 row : list
247 List of items in a single row.
248 control_row : list of bool
249 Helps to determine if particular column is in same parent index as the
250 previous value. Used to stop propagation of empty cells between
251 different indexes.
252
253 Returns
254 -------
255 Returns changed row and control_row
256 """
257 last = row[0]
258 for i in range(1, len(row)):
259 if not control_row[i]:
260 last = row[i]
261
262 if row[i] == "" or row[i] is None:
263 row[i] = last
264 else:
265 control_row[i] = False
266 last = row[i]
267
268 return row, control_row
269
270
271def pop_header_name(
272 row: list[Hashable], index_col: int | Sequence[int]
273) -> tuple[Hashable | None, list[Hashable]]:
274 """
275 Pop the header name for MultiIndex parsing.
276
277 Parameters
278 ----------
279 row : list
280 The data row to parse for the header name.
281 index_col : int, list
282 The index columns for our data. Assumed to be non-null.
283 If a list is passed, the maximum ``index_col`` value is used.
284
285 Returns
286 -------
287 header_name : str
288 The extracted header name.
289 trimmed_row : list
290 The original data row with the header name removed.
291 """
292 # Pop out header name and fill w/blank.
293 if is_list_like(index_col):
294 assert isinstance(index_col, Iterable)
295 i = max(index_col)
296 else:
297 assert not isinstance(index_col, Iterable)
298 i = index_col
299
300 header_name = row[i]
301 header_name = None if header_name == "" else header_name
302
303 return header_name, [*row[:i], "", *row[i + 1 :]]
304
305
306def combine_kwargs(
307 engine_kwargs: dict[str, Any] | None, kwargs: dict[str, Any]
308) -> dict[str, Any]:
309 """
310 Used to combine two sources of kwargs for the backend engine.
311
312 Use of kwargs is deprecated, this function is solely for use in 1.3 and should
313 be removed in 1.4/2.0. Also _base.ExcelWriter.__new__ ensures either engine_kwargs
314 or kwargs must be None or empty respectively.
315
316 Parameters
317 ----------
318 engine_kwargs: dict
319 kwargs to be passed through to the engine.
320 kwargs: dict
321 kwargs to be passed through to the engine (deprecated)
322
323 Returns
324 -------
325 engine_kwargs combined with kwargs
326 """
327 if engine_kwargs is None:
328 result = {}
329 else:
330 result = engine_kwargs.copy()
331 result.update(kwargs)
332 return result