@@ -600,12 +600,13 @@ def _cast_to_stata_types(data: DataFrame) -> DataFrame:
600
600
dtype = data [col ].dtype
601
601
for c_data in conversion_data :
602
602
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
606
607
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
609
610
if data [col ].max () >= 2 ** 53 :
610
611
ws = precision_loss_doc .format ("uint64" , "float64" )
611
612
@@ -967,17 +968,15 @@ def __init__(self) -> None:
967
968
(255 , np .dtype (np .float64 )),
968
969
]
969
970
)
970
- self .DTYPE_MAP_XML = {
971
+ self .DTYPE_MAP_XML : dict [ int , np . dtype ] = {
971
972
32768 : np .dtype (np .uint8 ), # Keys to GSO
972
973
65526 : np .dtype (np .float64 ),
973
974
65527 : np .dtype (np .float32 ),
974
975
65528 : np .dtype (np .int32 ),
975
976
65529 : np .dtype (np .int16 ),
976
977
65530 : np .dtype (np .int8 ),
977
978
}
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" ))
981
980
self .TYPE_MAP_XML = {
982
981
# Not really a Q, unclear how to handle byteswap
983
982
32768 : "Q" ,
@@ -1296,9 +1295,7 @@ def g(typ: int) -> str | np.dtype:
1296
1295
if typ <= 2045 :
1297
1296
return str (typ )
1298
1297
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 ]
1302
1299
except KeyError as err :
1303
1300
raise ValueError (f"cannot convert stata dtype [{ typ } ]" ) from err
1304
1301
0 commit comments