Skip to content

Commit 1346995

Browse files
committed
TYP: cols property
Limitations: - ignore type[assignment] error. - Created additional method _refine_cols to allow conversion from Optional[Sequence[Label]] to Sequence[Label].
1 parent f1e1ac8 commit 1346995

File tree

1 file changed

+9
-7
lines changed

1 file changed

+9
-7
lines changed

pandas/io/formats/csvs.py

+9-7
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def __init__(
111111
self.escapechar = escapechar
112112
self.line_terminator = line_terminator or os.linesep
113113
self.date_format = date_format
114-
self.cols = cols
114+
self.cols = cols # type: ignore[assignment]
115115
self.chunksize = chunksize # type: ignore[assignment]
116116

117117
@property
@@ -157,11 +157,14 @@ def has_mi_columns(self) -> bool:
157157
return bool(isinstance(self.obj.columns, ABCMultiIndex))
158158

159159
@property
160-
def cols(self):
160+
def cols(self) -> Sequence[Label]:
161161
return self._cols
162162

163163
@cols.setter
164-
def cols(self, cols):
164+
def cols(self, cols: Optional[Sequence[Label]]) -> None:
165+
self._cols = self._refine_cols(cols)
166+
167+
def _refine_cols(self, cols: Optional[Sequence[Label]]) -> Sequence[Label]:
165168
# validate mi options
166169
if self.has_mi_columns:
167170
if cols is not None:
@@ -179,11 +182,10 @@ def cols(self, cols):
179182
# and make sure sure cols is just a list of labels
180183
cols = self.obj.columns
181184
if isinstance(cols, ABCIndexClass):
182-
cols = cols.to_native_types(**self._number_format)
185+
return cols.to_native_types(**self._number_format)
183186
else:
184-
cols = list(cols)
185-
186-
self._cols = cols
187+
assert isinstance(cols, Sequence)
188+
return list(cols)
187189

188190
@property
189191
def _number_format(self) -> dict:

0 commit comments

Comments
 (0)