|
6 | 6 |
|
7 | 7 | """
|
8 | 8 | import operator
|
| 9 | +from typing import List, Set |
9 | 10 | import warnings
|
10 | 11 |
|
11 | 12 | import numpy as np
|
|
21 | 22 | import numexpr as ne
|
22 | 23 |
|
23 | 24 | _TEST_MODE = None
|
24 |
| -_TEST_RESULT = None |
| 25 | +_TEST_RESULT: List[bool] = list() |
25 | 26 | _USE_NUMEXPR = _NUMEXPR_INSTALLED
|
26 | 27 | _evaluate = None
|
27 | 28 | _where = None
|
@@ -75,7 +76,7 @@ def _can_use_numexpr(op, op_str, a, b, dtype_check):
|
75 | 76 | # required min elements (otherwise we are adding overhead)
|
76 | 77 | if np.prod(a.shape) > _MIN_ELEMENTS:
|
77 | 78 | # check for dtype compatibility
|
78 |
| - dtypes = set() |
| 79 | + dtypes: Set[str] = set() |
79 | 80 | for o in [a, b]:
|
80 | 81 | # Series implements dtypes, check for dimension count as well
|
81 | 82 | if hasattr(o, "dtypes") and o.ndim > 1:
|
@@ -247,25 +248,28 @@ def where(cond, a, b, use_numexpr=True):
|
247 | 248 | return _where(cond, a, b) if use_numexpr else _where_standard(cond, a, b)
|
248 | 249 |
|
249 | 250 |
|
250 |
| -def set_test_mode(v=True): |
| 251 | +def set_test_mode(v: bool = True) -> None: |
251 | 252 | """
|
252 |
| - Keeps track of whether numexpr was used. Stores an additional ``True`` |
253 |
| - for every successful use of evaluate with numexpr since the last |
254 |
| - ``get_test_result`` |
| 253 | + Keeps track of whether numexpr was used. |
| 254 | +
|
| 255 | + Stores an additional ``True`` for every successful use of evaluate with |
| 256 | + numexpr since the last ``get_test_result``. |
255 | 257 | """
|
256 | 258 | global _TEST_MODE, _TEST_RESULT
|
257 | 259 | _TEST_MODE = v
|
258 | 260 | _TEST_RESULT = []
|
259 | 261 |
|
260 | 262 |
|
261 |
| -def _store_test_result(used_numexpr): |
| 263 | +def _store_test_result(used_numexpr: bool) -> None: |
262 | 264 | global _TEST_RESULT
|
263 | 265 | if used_numexpr:
|
264 | 266 | _TEST_RESULT.append(used_numexpr)
|
265 | 267 |
|
266 | 268 |
|
267 |
| -def get_test_result(): |
268 |
| - """get test result and reset test_results""" |
| 269 | +def get_test_result() -> List[bool]: |
| 270 | + """ |
| 271 | + Get test result and reset test_results. |
| 272 | + """ |
269 | 273 | global _TEST_RESULT
|
270 | 274 | res = _TEST_RESULT
|
271 | 275 | _TEST_RESULT = []
|
|
0 commit comments