Skip to content

Commit 55de270

Browse files
author
Mateusz
committed
TYP: Fix io stata type ignores
GH37715
1 parent 6da08eb commit 55de270

File tree

1 file changed

+9
-12
lines changed

1 file changed

+9
-12
lines changed

pandas/io/stata.py

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -600,12 +600,13 @@ def _cast_to_stata_types(data: DataFrame) -> DataFrame:
600600
dtype = data[col].dtype
601601
for c_data in conversion_data:
602602
if dtype == c_data[0]:
603-
# Value of type variable "_IntType" of "iinfo" cannot be "object"
604-
if data[col].max() <= np.iinfo(c_data[1]).max: # type: ignore[type-var]
605-
dtype = c_data[1]
603+
c_data_small: type[np.signedinteger] = c_data[1]
604+
c_data_large: type[np.signedinteger | np.float64] = c_data[2]
605+
if data[col].max() <= np.iinfo(c_data_small).max:
606+
dtype = c_data_small
606607
else:
607-
dtype = c_data[2]
608-
if c_data[2] == np.int64: # Warn if necessary
608+
dtype = c_data_large
609+
if c_data_large == np.int64: # Warn if necessary
609610
if data[col].max() >= 2**53:
610611
ws = precision_loss_doc.format("uint64", "float64")
611612

@@ -967,17 +968,15 @@ def __init__(self) -> None:
967968
(255, np.dtype(np.float64)),
968969
]
969970
)
970-
self.DTYPE_MAP_XML = {
971+
self.DTYPE_MAP_XML: dict[int, np.dtype] = {
971972
32768: np.dtype(np.uint8), # Keys to GSO
972973
65526: np.dtype(np.float64),
973974
65527: np.dtype(np.float32),
974975
65528: np.dtype(np.int32),
975976
65529: np.dtype(np.int16),
976977
65530: np.dtype(np.int8),
977978
}
978-
# error: Argument 1 to "list" has incompatible type "str";
979-
# expected "Iterable[int]" [arg-type]
980-
self.TYPE_MAP = list(range(251)) + list("bhlfd") # type: ignore[arg-type]
979+
self.TYPE_MAP = list(tuple(range(251)) + tuple("bhlfd"))
981980
self.TYPE_MAP_XML = {
982981
# Not really a Q, unclear how to handle byteswap
983982
32768: "Q",
@@ -1296,9 +1295,7 @@ def g(typ: int) -> str | np.dtype:
12961295
if typ <= 2045:
12971296
return str(typ)
12981297
try:
1299-
# error: Incompatible return value type (got "Type[number]", expected
1300-
# "Union[str, dtype]")
1301-
return self.DTYPE_MAP_XML[typ] # type: ignore[return-value]
1298+
return self.DTYPE_MAP_XML[typ]
13021299
except KeyError as err:
13031300
raise ValueError(f"cannot convert stata dtype [{typ}]") from err
13041301

0 commit comments

Comments
 (0)