@@ -49,44 +49,64 @@ def tree(request, leaf_size):
49
49
class TestIntervalTree (object ):
50
50
51
51
def test_get_loc (self , tree ):
52
- tm .assert_numpy_array_equal (tree .get_loc (1 ),
53
- np .array ([0 ], dtype = 'int64' ))
54
- tm .assert_numpy_array_equal (np .sort (tree .get_loc (2 )),
55
- np .array ([0 , 1 ], dtype = 'int64' ))
52
+ result = tree .get_loc (1 )
53
+ expected = np .array ([0 ], dtype = 'intp' )
54
+ tm .assert_numpy_array_equal (result , expected )
55
+
56
+ result = np .sort (tree .get_loc (2 ))
57
+ expected = np .array ([0 , 1 ], dtype = 'intp' )
58
+ tm .assert_numpy_array_equal (result , expected )
59
+
56
60
with pytest .raises (KeyError ):
57
61
tree .get_loc (- 1 )
58
62
59
63
def test_get_indexer (self , tree ):
60
- tm .assert_numpy_array_equal (
61
- tree .get_indexer (np .array ([1.0 , 5.5 , 6.5 ])),
62
- np .array ([0 , 4 , - 1 ], dtype = 'int64' ))
64
+ result = tree .get_indexer (np .array ([1.0 , 5.5 , 6.5 ]))
65
+ expected = np .array ([0 , 4 , - 1 ], dtype = 'intp' )
66
+ tm .assert_numpy_array_equal (result , expected )
67
+
63
68
with pytest .raises (KeyError ):
64
69
tree .get_indexer (np .array ([3.0 ]))
65
70
66
71
def test_get_indexer_non_unique (self , tree ):
67
72
indexer , missing = tree .get_indexer_non_unique (
68
73
np .array ([1.0 , 2.0 , 6.5 ]))
69
- tm .assert_numpy_array_equal (indexer [:1 ],
70
- np .array ([0 ], dtype = 'int64' ))
71
- tm .assert_numpy_array_equal (np .sort (indexer [1 :3 ]),
72
- np .array ([0 , 1 ], dtype = 'int64' ))
73
- tm .assert_numpy_array_equal (np .sort (indexer [3 :]),
74
- np .array ([- 1 ], dtype = 'int64' ))
75
- tm .assert_numpy_array_equal (missing , np .array ([2 ], dtype = 'int64' ))
74
+
75
+ result = indexer [:1 ]
76
+ expected = np .array ([0 ], dtype = 'intp' )
77
+ tm .assert_numpy_array_equal (result , expected )
78
+
79
+ result = np .sort (indexer [1 :3 ])
80
+ expected = np .array ([0 , 1 ], dtype = 'intp' )
81
+ tm .assert_numpy_array_equal (result , expected )
82
+
83
+ result = np .sort (indexer [3 :])
84
+ expected = np .array ([- 1 ], dtype = 'intp' )
85
+ tm .assert_numpy_array_equal (result , expected )
86
+
87
+ result = missing
88
+ expected = np .array ([2 ], dtype = 'intp' )
89
+ tm .assert_numpy_array_equal (result , expected )
76
90
77
91
def test_duplicates (self , dtype ):
78
92
left = np .array ([0 , 0 , 0 ], dtype = dtype )
79
93
tree = IntervalTree (left , left + 1 )
80
- tm .assert_numpy_array_equal (np .sort (tree .get_loc (0.5 )),
81
- np .array ([0 , 1 , 2 ], dtype = 'int64' ))
94
+
95
+ result = np .sort (tree .get_loc (0.5 ))
96
+ expected = np .array ([0 , 1 , 2 ], dtype = 'intp' )
97
+ tm .assert_numpy_array_equal (result , expected )
82
98
83
99
with pytest .raises (KeyError ):
84
100
tree .get_indexer (np .array ([0.5 ]))
85
101
86
102
indexer , missing = tree .get_indexer_non_unique (np .array ([0.5 ]))
87
- tm .assert_numpy_array_equal (np .sort (indexer ),
88
- np .array ([0 , 1 , 2 ], dtype = 'int64' ))
89
- tm .assert_numpy_array_equal (missing , np .array ([], dtype = 'int64' ))
103
+ result = np .sort (indexer )
104
+ expected = np .array ([0 , 1 , 2 ], dtype = 'intp' )
105
+ tm .assert_numpy_array_equal (result , expected )
106
+
107
+ result = missing
108
+ expected = np .array ([], dtype = 'intp' )
109
+ tm .assert_numpy_array_equal (result , expected )
90
110
91
111
def test_get_loc_closed (self , closed ):
92
112
tree = IntervalTree ([0 ], [1 ], closed = closed )
@@ -96,8 +116,9 @@ def test_get_loc_closed(self, closed):
96
116
with pytest .raises (KeyError ):
97
117
tree .get_loc (p )
98
118
else :
99
- tm .assert_numpy_array_equal (tree .get_loc (p ),
100
- np .array ([0 ], dtype = 'int64' ))
119
+ result = tree .get_loc (p )
120
+ expected = np .array ([0 ], dtype = 'intp' )
121
+ tm .assert_numpy_array_equal (result , expected )
101
122
102
123
@pytest .mark .parametrize ('leaf_size' , [
103
124
skipif_32bit (1 ), skipif_32bit (10 ), skipif_32bit (100 ), 10000 ])
0 commit comments