|
1 | 1 | from __future__ import annotations
|
2 | 2 |
|
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 |
| - |
12 | 3 | import pandas as pd
|
13 | 4 | from pandas import DataFrame
|
| 5 | +from pandas._testing import ensure_clean |
14 | 6 | from typing_extensions import assert_type
|
15 | 7 |
|
16 | 8 | from tests import check
|
|
23 | 15 | DF = DataFrame({"a": [1, 2, 3], "b": [0.0, 0.0, 0.0]})
|
24 | 16 |
|
25 | 17 |
|
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 |
| - |
66 | 18 | def test_read_stata_df():
|
67 | 19 | with ensure_clean() as path:
|
68 | 20 | DF.to_stata(path)
|
69 | 21 | check(assert_type(read_stata(path), pd.DataFrame), pd.DataFrame)
|
70 | 22 |
|
71 | 23 |
|
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 |
| - |
87 | 24 | def test_read_stata_iterator():
|
88 | 25 | with ensure_clean() as path:
|
89 | 26 | str_path = str(path)
|
|
0 commit comments