1from collections import OrderedDict
2
3import numpy as np
4import pandas as pd
5import pytest
6
7from statsmodels.tools.validation import (
8    array_like,
9    PandasWrapper,
10    bool_like,
11    dict_like,
12    float_like,
13    int_like,
14    string_like,
15)
16
17from statsmodels.tools.validation.validation import _right_squeeze
18
19
20@pytest.fixture(params=[True, False])
21def use_pandas(request):
22    return request.param
23
24
25def gen_data(dim, use_pandas):
26    if dim == 1:
27        out = np.empty(10,)
28        if use_pandas:
29            out = pd.Series(out)
30    elif dim == 2:
31        out = np.empty((20, 10))
32        if use_pandas:
33            out = pd.DataFrame(out)
34    else:
35        out = np.empty(np.arange(5, 5 + dim))
36
37    return out
38
39
40class TestArrayLike(object):
41    def test_1d(self, use_pandas):
42        data = gen_data(1, use_pandas)
43        a = array_like(data, "a")
44        assert a.ndim == 1
45        assert a.shape == (10,)
46        assert type(a) is np.ndarray
47
48        a = array_like(data, "a", ndim=1)
49        assert a.ndim == 1
50        a = array_like(data, "a", shape=(10,))
51        assert a.shape == (10,)
52        a = array_like(data, "a", ndim=1, shape=(None,))
53        assert a.ndim == 1
54        a = array_like(data, "a", ndim=2, shape=(10, 1))
55        assert a.ndim == 2
56        assert a.shape == (10, 1)
57
58        with pytest.raises(ValueError, match="a is required to have shape"):
59            array_like(data, "a", shape=(5,))
60
61    def test_2d(self, use_pandas):
62        data = gen_data(2, use_pandas)
63        a = array_like(data, "a", ndim=2)
64        assert a.ndim == 2
65        assert a.shape == (20, 10)
66        assert type(a) is np.ndarray
67
68        a = array_like(data, "a", ndim=2)
69        assert a.ndim == 2
70        a = array_like(data, "a", ndim=2, shape=(20, None))
71        assert a.shape == (20, 10)
72        a = array_like(data, "a", ndim=2, shape=(20,))
73        assert a.shape == (20, 10)
74        a = array_like(data, "a", ndim=2, shape=(None, 10))
75        assert a.shape == (20, 10)
76
77        a = array_like(data, "a", ndim=2, shape=(None, None))
78        assert a.ndim == 2
79        a = array_like(data, "a", ndim=3)
80        assert a.ndim == 3
81        assert a.shape == (20, 10, 1)
82
83        with pytest.raises(ValueError, match="a is required to have shape"):
84            array_like(data, "a", ndim=2, shape=(10,))
85        with pytest.raises(ValueError, match="a is required to have shape"):
86            array_like(data, "a", ndim=2, shape=(20, 20))
87        with pytest.raises(ValueError, match="a is required to have shape"):
88            array_like(data, "a", ndim=2, shape=(None, 20))
89        match = "a is required to have ndim 1 but has ndim 2"
90        with pytest.raises(ValueError, match=match):
91            array_like(data, "a", ndim=1)
92        match = "a must have ndim <= 1"
93        with pytest.raises(ValueError, match=match):
94            array_like(data, "a", maxdim=1)
95
96    def test_3d(self):
97        data = gen_data(3, False)
98        a = array_like(data, "a", ndim=3)
99        assert a.shape == (5, 6, 7)
100        assert a.ndim == 3
101        assert type(a) is np.ndarray
102
103        a = array_like(data, "a", ndim=3, shape=(5, None, 7))
104        assert a.shape == (5, 6, 7)
105        a = array_like(data, "a", ndim=3, shape=(None, None, 7))
106        assert a.shape == (5, 6, 7)
107        a = array_like(data, "a", ndim=5)
108        assert a.shape == (5, 6, 7, 1, 1)
109        with pytest.raises(ValueError, match="a is required to have shape"):
110            array_like(data, "a", ndim=3, shape=(10,))
111        with pytest.raises(ValueError, match="a is required to have shape"):
112            array_like(data, "a", ndim=3, shape=(None, None, 5))
113        match = "a is required to have ndim 2 but has ndim 3"
114        with pytest.raises(ValueError, match=match):
115            array_like(data, "a", ndim=2)
116        match = "a must have ndim <= 1"
117        with pytest.raises(ValueError, match=match):
118            array_like(data, "a", maxdim=1)
119        match = "a must have ndim <= 2"
120        with pytest.raises(ValueError, match=match):
121            array_like(data, "a", maxdim=2)
122
123    def test_right_squeeze_and_pad(self):
124        data = np.empty((2, 1, 2))
125        a = array_like(data, "a", ndim=3)
126        assert a.shape == (2, 1, 2)
127        data = np.empty((2))
128        a = array_like(data, "a", ndim=3)
129        assert a.shape == (2, 1, 1)
130        data = np.empty((2, 1))
131        a = array_like(data, "a", ndim=3)
132        assert a.shape == (2, 1, 1)
133
134        data = np.empty((2, 1, 1, 1))
135        a = array_like(data, "a", ndim=3)
136        assert a.shape == (2, 1, 1)
137
138        data = np.empty((2, 1, 1, 2, 1, 1))
139        with pytest.raises(ValueError):
140            array_like(data, "a", ndim=3)
141
142    def test_contiguous(self):
143        x = np.arange(10)
144        y = x[::2]
145        a = array_like(y, "a", contiguous=True)
146        assert not y.flags["C_CONTIGUOUS"]
147        assert a.flags["C_CONTIGUOUS"]
148
149    def test_dtype(self):
150        x = np.arange(10)
151        a = array_like(x, "a", dtype=np.float32)
152        assert a.dtype == np.float32
153
154        a = array_like(x, "a", dtype=np.uint8)
155        assert a.dtype == np.uint8
156
157    @pytest.mark.xfail(reason="Failing for now")
158    def test_dot(self, use_pandas):
159        data = gen_data(2, use_pandas)
160        a = array_like(data, "a")
161        assert not isinstance(a.T.dot(data), array_like)
162        assert not isinstance(a.T.dot(a), array_like)
163
164    def test_slice(self, use_pandas):
165        data = gen_data(2, use_pandas)
166        a = array_like(data, "a", ndim=2)
167        assert type(a[1:]) is np.ndarray
168
169
170def test_right_squeeze():
171    x = np.empty((10, 1, 10))
172    y = _right_squeeze(x)
173    assert y.shape == (10, 1, 10)
174
175    x = np.empty((10, 10, 1))
176    y = _right_squeeze(x)
177    assert y.shape == (10, 10)
178
179    x = np.empty((10, 10, 1, 1, 1, 1, 1))
180    y = _right_squeeze(x)
181    assert y.shape == (10, 10)
182
183    x = np.empty((10, 1, 10, 1, 1, 1, 1, 1))
184    y = _right_squeeze(x)
185    assert y.shape == (10, 1, 10)
186
187
188def test_wrap_pandas(use_pandas):
189    a = gen_data(1, use_pandas)
190    b = gen_data(1, False)
191
192    wrapped = PandasWrapper(a).wrap(b)
193    expected_type = pd.Series if use_pandas else np.ndarray
194    assert isinstance(wrapped, expected_type)
195    assert not use_pandas or wrapped.name is None
196
197    wrapped = PandasWrapper(a).wrap(b, columns="name")
198    assert isinstance(wrapped, expected_type)
199    assert not use_pandas or wrapped.name == "name"
200
201    wrapped = PandasWrapper(a).wrap(b, columns=["name"])
202    assert isinstance(wrapped, expected_type)
203    assert not use_pandas or wrapped.name == "name"
204
205    expected_type = pd.DataFrame if use_pandas else np.ndarray
206    wrapped = PandasWrapper(a).wrap(b[:, None])
207    assert isinstance(wrapped, expected_type)
208    assert not use_pandas or wrapped.columns[0] == 0
209
210    wrapped = PandasWrapper(a).wrap(b[:, None], columns=["name"])
211    assert isinstance(wrapped, expected_type)
212    assert not use_pandas or wrapped.columns == ["name"]
213
214    if use_pandas:
215        match = "Can only wrap 1 or 2-d array_like"
216        with pytest.raises(ValueError, match=match):
217            PandasWrapper(a).wrap(b[:, None, None])
218
219        match = "obj must have the same number of elements in axis 0 as"
220        with pytest.raises(ValueError, match=match):
221            PandasWrapper(a).wrap(b[: b.shape[0] // 2])
222
223
224def test_wrap_pandas_append():
225    a = gen_data(1, True)
226    a.name = "apple"
227    b = gen_data(1, False)
228    wrapped = PandasWrapper(a).wrap(b, append="appended")
229    expected = "apple_appended"
230    assert wrapped.name == expected
231
232    a = gen_data(2, True)
233    a.columns = ["apple_" + str(i) for i in range(a.shape[1])]
234    b = gen_data(2, False)
235    wrapped = PandasWrapper(a).wrap(b, append="appended")
236    expected = [c + "_appended" for c in a.columns]
237    assert list(wrapped.columns) == expected
238
239
240def test_wrap_pandas_append_non_string():
241    # GH 6826
242    a = gen_data(1, True)
243    a.name = 7
244    b = gen_data(1, False)
245    wrapped = PandasWrapper(a).wrap(b, append="appended")
246    expected = "7_appended"
247    assert wrapped.name == expected
248
249    a = gen_data(2, True)
250    a.columns = [i for i in range(a.shape[1])]
251    b = gen_data(2, False)
252    wrapped = PandasWrapper(a).wrap(b, append="appended")
253    expected = [f"{c}_appended" for c in a.columns]
254    assert list(wrapped.columns) == expected
255
256
257class CustomDict(dict):
258    pass
259
260
261@pytest.fixture(params=(dict, OrderedDict, CustomDict, None))
262def dict_type(request):
263    return request.param
264
265
266def test_optional_dict_like(dict_type):
267    val = dict_type() if dict_type is not None else dict_type
268    out = dict_like(val, "value", optional=True)
269    assert isinstance(out, type(val))
270
271
272def test_optional_dict_like_error():
273    match = r"value must be a dict or dict_like \(i.e., a Mapping\)"
274    with pytest.raises(TypeError, match=match):
275        dict_like([], "value", optional=True)
276    with pytest.raises(TypeError, match=match):
277        dict_like({"a"}, "value", optional=True)
278    with pytest.raises(TypeError, match=match):
279        dict_like("a", "value", optional=True)
280
281
282def test_string():
283    out = string_like("apple", "value")
284    assert out == "apple"
285
286    out = string_like("apple", "value", options=("apple", "banana", "cherry"))
287    assert out == "apple"
288
289    with pytest.raises(TypeError, match="value must be a string"):
290        string_like(1, "value")
291    with pytest.raises(TypeError, match="value must be a string"):
292        string_like(b"4", "value")
293    with pytest.raises(
294        ValueError,
295        match="value must be one of: 'apple'," " 'banana', 'cherry'",
296    ):
297        string_like("date", "value", options=("apple", "banana", "cherry"))
298
299
300def test_optional_string():
301    out = string_like("apple", "value")
302    assert out == "apple"
303
304    out = string_like("apple", "value", options=("apple", "banana", "cherry"))
305    assert out == "apple"
306
307    out = string_like(None, "value", optional=True)
308    assert out is None
309
310    out = string_like(
311        None, "value", optional=True, options=("apple", "banana", "cherry")
312    )
313    assert out is None
314
315    with pytest.raises(TypeError, match="value must be a string"):
316        string_like(1, "value", optional=True)
317    with pytest.raises(TypeError, match="value must be a string"):
318        string_like(b"4", "value", optional=True)
319
320
321@pytest.fixture(params=(1.0, 1.1, np.float32(1.2), np.array([1.2]), 1.2 + 0j))
322def floating(request):
323    return request.param
324
325
326@pytest.fixture(params=(np.empty(2), 1.2 + 1j, True, "3.2", None))
327def not_floating(request):
328    return request.param
329
330
331def test_float_like(floating):
332    assert isinstance(float_like(floating, "floating"), float)
333    assert isinstance(float_like(floating, "floating", optional=True), float)
334    assert float_like(None, "floating", optional=True) is None
335    if isinstance(floating, (int, np.integer, float, np.inexact)):
336        assert isinstance(float_like(floating, "floating", strict=True), float)
337        assert float_like(None, "floating", optional=True, strict=True) is None
338
339
340def test_not_float_like(not_floating):
341    with pytest.raises(TypeError):
342        float_like(not_floating, "floating")
343
344
345@pytest.fixture(params=(1.0, 2, np.float32(3.0), np.array([4.0])))
346def integer(request):
347    return request.param
348
349
350@pytest.fixture(
351    params=(
352        3.2,
353        np.float32(3.2),
354        3 + 2j,
355        complex(2.3 + 0j),
356        "apple",
357        1.0 + 0j,
358        np.timedelta64(2),
359    )
360)
361def not_integer(request):
362    return request.param
363
364
365def test_int_like(integer):
366    assert isinstance(int_like(integer, "integer"), int)
367    assert isinstance(int_like(integer, "integer", optional=True), int)
368    assert int_like(None, "floating", optional=True) is None
369    if isinstance(integer, (int, np.integer)):
370        assert isinstance(int_like(integer, "integer", strict=True), int)
371        assert int_like(None, "floating", optional=True, strict=True) is None
372
373
374def test_not_int_like(not_integer):
375    with pytest.raises(TypeError):
376        int_like(not_integer, "integer")
377
378
379@pytest.fixture(params=[True, False, 1, 1.2, "a", ""])
380def boolean(request):
381    return request.param
382
383
384def test_bool_like(boolean):
385    assert isinstance(bool_like(boolean, "boolean"), bool)
386    assert bool_like(None, "boolean", optional=True) is None
387    if isinstance(boolean, bool):
388        assert isinstance(bool_like(boolean, "boolean", strict=True), bool)
389    else:
390        with pytest.raises(TypeError):
391            bool_like(boolean, "boolean", strict=True)
392
393
394def test_not_bool_like():
395    with pytest.raises(TypeError):
396        bool_like(np.array([True, True]), boolean)
397