Skip to content

Commit 9f5e609

Browse files
authored
Type read_fwf (#44678)
1 parent f5107e4 commit 9f5e609

File tree

2 files changed

+30
-13
lines changed

2 files changed

+30
-13
lines changed

pandas/io/parsers/python_parser.py

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,11 @@
99
import re
1010
import sys
1111
from typing import (
12+
IO,
1213
DefaultDict,
1314
Hashable,
1415
Iterator,
16+
Literal,
1517
Mapping,
1618
Sequence,
1719
cast,
@@ -1135,9 +1137,17 @@ class FixedWidthReader(abc.Iterator):
11351137
A reader of fixed-width lines.
11361138
"""
11371139

1138-
def __init__(self, f, colspecs, delimiter, comment, skiprows=None, infer_nrows=100):
1140+
def __init__(
1141+
self,
1142+
f: IO[str],
1143+
colspecs: list[tuple[int, int]] | Literal["infer"],
1144+
delimiter: str | None,
1145+
comment: str | None,
1146+
skiprows: set[int] | None = None,
1147+
infer_nrows: int = 100,
1148+
) -> None:
11391149
self.f = f
1140-
self.buffer = None
1150+
self.buffer: Iterator | None = None
11411151
self.delimiter = "\r\n" + delimiter if delimiter else "\n\r\t "
11421152
self.comment = comment
11431153
if colspecs == "infer":
@@ -1165,7 +1175,7 @@ def __init__(self, f, colspecs, delimiter, comment, skiprows=None, infer_nrows=1
11651175
"2 element tuple or list of integers"
11661176
)
11671177

1168-
def get_rows(self, infer_nrows, skiprows=None):
1178+
def get_rows(self, infer_nrows: int, skiprows: set[int] | None = None) -> list[str]:
11691179
"""
11701180
Read rows from self.f, skipping as specified.
11711181
@@ -1203,7 +1213,9 @@ def get_rows(self, infer_nrows, skiprows=None):
12031213
self.buffer = iter(buffer_rows)
12041214
return detect_rows
12051215

1206-
def detect_colspecs(self, infer_nrows=100, skiprows=None):
1216+
def detect_colspecs(
1217+
self, infer_nrows: int = 100, skiprows: set[int] | None = None
1218+
) -> list[tuple[int, int]]:
12071219
# Regex escape the delimiters
12081220
delimiters = "".join([fr"\{x}" for x in self.delimiter])
12091221
pattern = re.compile(f"([^{delimiters}]+)")
@@ -1223,7 +1235,7 @@ def detect_colspecs(self, infer_nrows=100, skiprows=None):
12231235
edge_pairs = list(zip(edges[::2], edges[1::2]))
12241236
return edge_pairs
12251237

1226-
def __next__(self):
1238+
def __next__(self) -> list[str]:
12271239
if self.buffer is not None:
12281240
try:
12291241
line = next(self.buffer)
@@ -1242,13 +1254,15 @@ class FixedWidthFieldParser(PythonParser):
12421254
See PythonParser for details.
12431255
"""
12441256

1245-
def __init__(self, f, **kwds):
1257+
def __init__(
1258+
self, f: FilePath | ReadCsvBuffer[bytes] | ReadCsvBuffer[str], **kwds
1259+
) -> None:
12461260
# Support iterators, convert to a list.
12471261
self.colspecs = kwds.pop("colspecs")
12481262
self.infer_nrows = kwds.pop("infer_nrows")
12491263
PythonParser.__init__(self, f, **kwds)
12501264

1251-
def _make_reader(self, f):
1265+
def _make_reader(self, f: IO[str]) -> None:
12521266
self.data = FixedWidthReader(
12531267
f,
12541268
self.colspecs,
@@ -1258,7 +1272,7 @@ def _make_reader(self, f):
12581272
self.infer_nrows,
12591273
)
12601274

1261-
def _remove_empty_lines(self, lines) -> list:
1275+
def _remove_empty_lines(self, lines: list[list[Scalar]]) -> list[list[Scalar]]:
12621276
"""
12631277
Returns the list of lines without the empty ones. With fixed-width
12641278
fields, empty lines become arrays of empty strings.

pandas/io/parsers/readers.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -760,11 +760,11 @@ def read_table(
760760

761761
def read_fwf(
762762
filepath_or_buffer: FilePath | ReadCsvBuffer[bytes] | ReadCsvBuffer[str],
763-
colspecs="infer",
764-
widths=None,
765-
infer_nrows=100,
763+
colspecs: list[tuple[int, int]] | str | None = "infer",
764+
widths: list[int] | None = None,
765+
infer_nrows: int = 100,
766766
**kwds,
767-
):
767+
) -> DataFrame | TextFileReader:
768768
r"""
769769
Read a table of fixed-width formatted lines into DataFrame.
770770
@@ -799,7 +799,7 @@ def read_fwf(
799799
800800
Returns
801801
-------
802-
DataFrame or TextParser
802+
DataFrame or TextFileReader
803803
A comma-separated values (csv) file is returned as two-dimensional
804804
data structure with labeled axes.
805805
@@ -825,6 +825,9 @@ def read_fwf(
825825
colspecs.append((col, col + w))
826826
col += w
827827

828+
# for mypy
829+
assert colspecs is not None
830+
828831
# GH#40830
829832
# Ensure length of `colspecs` matches length of `names`
830833
names = kwds.get("names")

0 commit comments

Comments
 (0)