Skip to content

Commit 72c7fbf

Browse files
committed
CLN: Use ensure_clean from pandas
Use ensure clean Switch to only 1.5 compatible typing
1 parent 620f297 commit 72c7fbf

File tree

2 files changed

+3
-80
lines changed

2 files changed

+3
-80
lines changed

pandas-stubs/io/stata.pyi

Lines changed: 2 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ from pandas._typing import (
2727
@overload
2828
def read_stata(
2929
path: FilePath | ReadBuffer[bytes],
30+
*,
3031
convert_dates: bool = ...,
3132
convert_categoricals: bool = ...,
3233
index_col: str | None = ...,
@@ -35,29 +36,14 @@ def read_stata(
3536
columns: list[HashableT] | None = ...,
3637
order_categoricals: bool = ...,
3738
chunksize: int | None = ...,
38-
*,
39-
iterator: Literal[True],
40-
compression: CompressionOptions = ...,
41-
storage_options: StorageOptions = ...,
42-
) -> StataReader: ...
43-
@overload
44-
def read_stata(
45-
path: FilePath | ReadBuffer[bytes],
46-
convert_dates: bool,
47-
convert_categoricals: bool,
48-
index_col: str | None,
49-
convert_missing: bool,
50-
preserve_dtypes: bool,
51-
columns: list[HashableT] | None,
52-
order_categoricals: bool,
53-
chunksize: int | None,
5439
iterator: Literal[True],
5540
compression: CompressionOptions = ...,
5641
storage_options: StorageOptions = ...,
5742
) -> StataReader: ...
5843
@overload
5944
def read_stata(
6045
path: FilePath | ReadBuffer[bytes],
46+
*,
6147
convert_dates: bool = ...,
6248
convert_categoricals: bool = ...,
6349
index_col: str | None = ...,

tests/test_io.py

Lines changed: 1 addition & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,8 @@
11
from __future__ import annotations
22

3-
from contextlib import contextmanager
4-
from pathlib import Path
5-
import tempfile
6-
from typing import (
7-
IO,
8-
Any,
9-
)
10-
import uuid
11-
123
import pandas as pd
134
from pandas import DataFrame
5+
from pandas._testing import ensure_clean
146
from typing_extensions import assert_type
157

168
from tests import check
@@ -23,67 +15,12 @@
2315
DF = DataFrame({"a": [1, 2, 3], "b": [0.0, 0.0, 0.0]})
2416

2517

26-
@contextmanager
27-
def ensure_clean(filename=None, return_filelike: bool = False, **kwargs: Any):
28-
"""
29-
Gets a temporary path and agrees to remove on close.
30-
This implementation does not use tempfile.mkstemp to avoid having a file handle.
31-
If the code using the returned path wants to delete the file itself, windows
32-
requires that no program has a file handle to it.
33-
Parameters
34-
----------
35-
filename : str (optional)
36-
suffix of the created file.
37-
return_filelike : bool (default False)
38-
if True, returns a file-like which is *always* cleaned. Necessary for
39-
savefig and other functions which want to append extensions.
40-
**kwargs
41-
Additional keywords are passed to open().
42-
"""
43-
folder = Path(tempfile.gettempdir())
44-
45-
if filename is None:
46-
filename = ""
47-
filename = str(uuid.uuid4()) + filename
48-
path = folder / filename
49-
50-
path.touch()
51-
52-
handle_or_str: str | IO = str(path)
53-
if return_filelike:
54-
kwargs.setdefault("mode", "w+b")
55-
handle_or_str = open(path, **kwargs)
56-
57-
try:
58-
yield handle_or_str
59-
finally:
60-
if not isinstance(handle_or_str, str):
61-
handle_or_str.close()
62-
if path.is_file():
63-
path.unlink()
64-
65-
6618
def test_read_stata_df():
6719
with ensure_clean() as path:
6820
DF.to_stata(path)
6921
check(assert_type(read_stata(path), pd.DataFrame), pd.DataFrame)
7022

7123

72-
def test_read_stata_iterator_positional():
73-
with ensure_clean() as path:
74-
str_path = str(path)
75-
DF.to_stata(str_path)
76-
check(
77-
assert_type(
78-
read_stata(
79-
str_path, False, False, None, False, False, None, False, 2, True
80-
),
81-
StataReader,
82-
),
83-
StataReader,
84-
)
85-
86-
8724
def test_read_stata_iterator():
8825
with ensure_clean() as path:
8926
str_path = str(path)

0 commit comments

Comments
 (0)