Skip to content

Commit 4890dba

Browse files
bashtageKevin Sheppard
and
Kevin Sheppard
authored
ENH: Improve to_string typing (#244)
* ENH: Improve to_string typing * BUG/TST: Corect types and add test * MAINT: Remove dep on matplotlib * MAINT: Use literal and add Series Co-authored-by: Kevin Sheppard <kevin.sheppard@gmail.com>
1 parent b072842 commit 4890dba

File tree

6 files changed

+44
-21
lines changed

6 files changed

+44
-21
lines changed

pandas-stubs/_typing.pyi

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ from pandas._libs.tslibs import (
3939

4040
from pandas.core.dtypes.dtypes import ExtensionDtype
4141

42+
from pandas.io.formats.format import EngFormatter
43+
4244
ArrayLike = Union[ExtensionArray, np.ndarray]
4345
AnyArrayLike = Union[Index, Series, np.ndarray]
4446
PythonScalar = Union[str, bool, complex]
@@ -188,7 +190,10 @@ CompressionDict = dict[str, Any]
188190
CompressionOptions = Optional[
189191
Union[Literal["infer", "gzip", "bz2", "zip", "xz", "zstd"], CompressionDict]
190192
]
191-
193+
FormattersType = Union[
194+
list[Callable], tuple[Callable, ...], Mapping[Union[str, int], Callable]
195+
]
196+
FloatFormatType = str | Callable | EngFormatter
192197
# converters
193198
ConvertersArg = dict[Hashable, Callable[[Dtype], Dtype]]
194199

pandas-stubs/core/frame.pyi

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ from pandas._typing import (
5757
FilePath,
5858
FilePathOrBuffer,
5959
FillnaOptions,
60+
FloatFormatType,
61+
FormattersType,
6062
GroupByObjectNonScalar,
6163
HashableT,
6264
IgnoreRaise,
@@ -347,7 +349,7 @@ class DataFrame(NDFrame, OpsMixin):
347349
def to_html(
348350
self,
349351
buf: FilePath | WriteBuffer[str],
350-
columns: list[HashableT] | None = ...,
352+
columns: list[HashableT] | Index | Series | None = ...,
351353
col_space: ColspaceArgType | None = ...,
352354
header: _bool = ...,
353355
index: _bool = ...,
@@ -1928,7 +1930,7 @@ class DataFrame(NDFrame, OpsMixin):
19281930
mode: _str = ...,
19291931
encoding: _str | None = ...,
19301932
compression: _str | Mapping[_str, _str] = ...,
1931-
quoting: int | None = ...,
1933+
quoting: Literal[0, 1, 2, 3] | None = ...,
19321934
quotechar: _str = ...,
19331935
line_terminator: _str | None = ...,
19341936
chunksize: int | None = ...,
@@ -1952,7 +1954,7 @@ class DataFrame(NDFrame, OpsMixin):
19521954
mode: _str = ...,
19531955
encoding: _str | None = ...,
19541956
compression: _str | Mapping[_str, _str] = ...,
1955-
quoting: int | None = ...,
1957+
quoting: Literal[0, 1, 2, 3] | None = ...,
19561958
quotechar: _str = ...,
19571959
line_terminator: _str | None = ...,
19581960
chunksize: int | None = ...,
@@ -2003,45 +2005,46 @@ class DataFrame(NDFrame, OpsMixin):
20032005
@overload
20042006
def to_string(
20052007
self,
2006-
buf: FilePathOrBuffer | None,
2007-
columns: Sequence[_str] | None = ...,
2008-
col_space: int | list[int] | dict[_str | int, int] | None = ...,
2009-
header: _bool | Sequence[_str] = ...,
2008+
buf: FilePath | WriteBuffer[str],
2009+
columns: list[HashableT] | Index | Series | None = ...,
2010+
col_space: int | list[int] | dict[HashableT, int] | None = ...,
2011+
header: _bool | list[_str] | tuple[str, ...] = ...,
20102012
index: _bool = ...,
20112013
na_rep: _str = ...,
2012-
formatters=...,
2013-
float_format=...,
2014+
formatters: FormattersType | None = ...,
2015+
float_format: FloatFormatType | None = ...,
20142016
sparsify: _bool | None = ...,
20152017
index_names: _bool = ...,
20162018
justify: _str | None = ...,
20172019
max_rows: int | None = ...,
2018-
min_rows: int | None = ...,
20192020
max_cols: int | None = ...,
20202021
show_dimensions: _bool = ...,
20212022
decimal: _str = ...,
20222023
line_width: int | None = ...,
2024+
min_rows: int | None = ...,
20232025
max_colwidth: int | None = ...,
20242026
encoding: _str | None = ...,
20252027
) -> None: ...
20262028
@overload
20272029
def to_string(
20282030
self,
2029-
columns: Sequence[_str] | None = ...,
2030-
col_space: int | list[int] | dict[_str | int, int] | None = ...,
2031+
buf: None = ...,
2032+
columns: list[HashableT] | Index | Series | None = ...,
2033+
col_space: int | list[int] | dict[Hashable, int] | None = ...,
20312034
header: _bool | Sequence[_str] = ...,
20322035
index: _bool = ...,
20332036
na_rep: _str = ...,
2034-
formatters=...,
2035-
float_format=...,
2037+
formatters: FormattersType | None = ...,
2038+
float_format: FloatFormatType | None = ...,
20362039
sparsify: _bool | None = ...,
20372040
index_names: _bool = ...,
20382041
justify: _str | None = ...,
20392042
max_rows: int | None = ...,
2040-
min_rows: int | None = ...,
20412043
max_cols: int | None = ...,
20422044
show_dimensions: _bool = ...,
20432045
decimal: _str = ...,
20442046
line_width: int | None = ...,
2047+
min_rows: int | None = ...,
20452048
max_colwidth: int | None = ...,
20462049
encoding: _str | None = ...,
20472050
) -> _str: ...

pandas-stubs/core/generic.pyi

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -242,7 +242,7 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
242242
mode: _str = ...,
243243
encoding: _str | None = ...,
244244
compression: _str | Mapping[_str, _str] = ...,
245-
quoting: int | None = ...,
245+
quoting: Literal[0, 1, 2, 3] | None = ...,
246246
quotechar: _str = ...,
247247
line_terminator: _str | None = ...,
248248
chunksize: int | None = ...,
@@ -266,7 +266,7 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
266266
mode: _str = ...,
267267
encoding: _str | None = ...,
268268
compression: _str | Mapping[_str, _str] = ...,
269-
quoting: int | None = ...,
269+
quoting: Literal[0, 1, 2, 3] | None = ...,
270270
quotechar: _str = ...,
271271
line_terminator: _str | None = ...,
272272
chunksize: int | None = ...,

pandas-stubs/io/clipboards.pyi

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def read_clipboard(
7171
decimal: str = ...,
7272
lineterminator: str | None = ...,
7373
quotechar: str = ...,
74-
quoting: int = ...,
74+
quoting: Literal[0, 1, 2, 3] | None = ...,
7575
doublequote: bool = ...,
7676
escapechar: str | None = ...,
7777
comment: str | None = ...,
@@ -137,7 +137,7 @@ def read_clipboard(
137137
decimal: str = ...,
138138
lineterminator: str | None = ...,
139139
quotechar: str = ...,
140-
quoting: int = ...,
140+
quoting: Literal[0, 1, 2, 3] | None = ...,
141141
doublequote: bool = ...,
142142
escapechar: str | None = ...,
143143
comment: str | None = ...,
@@ -203,7 +203,7 @@ def read_clipboard(
203203
decimal: str = ...,
204204
lineterminator: str | None = ...,
205205
quotechar: str = ...,
206-
quoting: int = ...,
206+
quoting: Literal[0, 1, 2, 3] | None = ...,
207207
doublequote: bool = ...,
208208
escapechar: str | None = ...,
209209
comment: str | None = ...,

tests/test_frame.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
from collections import defaultdict
4+
import csv
45
import datetime
56
import io
67
from pathlib import Path
@@ -114,6 +115,9 @@ def test_types_to_csv() -> None:
114115
# Testing support for binary file handles, added in 1.2.0 https://pandas.pydata.org/docs/whatsnew/v1.2.0.html
115116
df.to_csv(io.BytesIO(), encoding="utf-8", compression="gzip")
116117

118+
# Testing support for binary file handles, added in 1.2.0 https://pandas.pydata.org/docs/whatsnew/v1.2.0.html
119+
df.to_csv(io.BytesIO(), quoting=csv.QUOTE_ALL, encoding="utf-8", compression="gzip")
120+
117121

118122
def test_types_to_csv_when_path_passed() -> None:
119123
df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})

tests/test_io.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,17 @@ def test_feather():
353353
check(assert_type(read_feather(bio), DataFrame), DataFrame)
354354

355355

356+
def test_to_string():
357+
check(assert_type(DF.to_string(), str), str)
358+
with ensure_clean() as path:
359+
check(assert_type(DF.to_string(path), None), type(None))
360+
check(assert_type(DF.to_string(pathlib.Path(path)), None), type(None))
361+
with open(path, "wt") as df_string:
362+
check(assert_type(DF.to_string(df_string), None), type(None))
363+
sio = io.StringIO()
364+
check(assert_type(DF.to_string(sio), None), type(None))
365+
366+
356367
def test_read_sql():
357368
with ensure_clean() as path:
358369
con = sqlite3.connect(path)

0 commit comments

Comments
 (0)