-
Notifications
You must be signed in to change notification settings - Fork 34
/
util.py
79 lines (69 loc) · 2.41 KB
/
util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#
# Copyright (c) 2020 IBM Corp.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#
# util.py
#
# Part of text_extensions_for_pandas
#
# Internal utility functions, not exposed in the public API.
#
from typing import Any
import numpy as np
from typing import *
import unittest
# Internal imports
_ELLIPSIS = " [...] "
_ELLIPSIS_LEN = len(_ELLIPSIS)
class TestBase(unittest.TestCase):
"""
Base class to hold common utility code used by test cases in multiple files.
"""
def _assertArrayEquals(self, a1: Union[np.ndarray, List[Any]],
a2: Union[np.ndarray, List[Any]]) -> None:
"""
Assert that two arrays are completely identical, with useful error
messages if they are not.
:param a1: first array to compare. Lists automatically converted to
arrays.
:param a2: second array (or list)
"""
a1 = np.array(a1) if isinstance(a1, np.ndarray) else a1
a2 = np.array(a2) if isinstance(a2, np.ndarray) else a2
if len(a1) != len(a2):
raise self.failureException(
f"Arrays:\n"
f" {a1}\n"
f"and\n"
f" {a2}\n"
f"have different lengths {len(a1)} and {len(a2)}"
)
mask = (a1 == a2)
if not np.all(mask):
raise self.failureException(
f"Arrays:\n"
f" {a1}\n"
f"and\n"
f" {a2}\n"
f"differ at positions: {np.argwhere(~mask)}"
)
def to_int_array(arr: Any) -> np.ndarray:
"""
Turn an input into a Numpy array with an integer dtype, avoiding making a copy
if possible.
"""
if isinstance(arr, np.ndarray) and np.issubdtype(arr.dtype, np.integer):
return arr # Avoid making a copy even if the input is an unusual integer dtype
else:
return np.array(arr, dtype=np.int32, copy=False)