forked from geopandas/geopandas
/
test_sindex.py
123 lines (100 loc) · 4.19 KB
/
test_sindex.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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import sys
from shapely.geometry import Point, Polygon
import geopandas
from geopandas import GeoDataFrame, GeoSeries, base, read_file
import pytest
@pytest.mark.skipif(sys.platform.startswith("win"), reason="fails on AppVeyor")
@pytest.mark.skipif(not base.HAS_SINDEX, reason="Rtree absent, skipping")
class TestSeriesSindex:
def test_empty_geoseries(self):
assert GeoSeries().sindex is None
def test_point(self):
s = GeoSeries([Point(0, 0)])
assert s.sindex.size == 1
hits = s.sindex.intersection((-1, -1, 1, 1))
assert len(list(hits)) == 1
hits = s.sindex.intersection((-2, -2, -1, -1))
assert len(list(hits)) == 0
def test_empty_point(self):
s = GeoSeries([Point()])
assert s.sindex is None
assert s._sindex_generated is True
def test_polygons(self):
t1 = Polygon([(0, 0), (1, 0), (1, 1)])
t2 = Polygon([(0, 0), (1, 1), (0, 1)])
sq = Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
s = GeoSeries([t1, t2, sq])
assert s.sindex.size == 3
def test_polygons_append(self):
t1 = Polygon([(0, 0), (1, 0), (1, 1)])
t2 = Polygon([(0, 0), (1, 1), (0, 1)])
sq = Polygon([(0, 0), (1, 0), (1, 1), (0, 1)])
s = GeoSeries([t1, t2, sq])
t = GeoSeries([t1, t2, sq], [3, 4, 5])
s = s.append(t)
assert len(s) == 6
assert s.sindex.size == 6
def test_lazy_build(self):
s = GeoSeries([Point(0, 0)])
assert s._sindex is None
assert s.sindex.size == 1
assert s._sindex is not None
@pytest.mark.skipif(sys.platform.startswith("win"), reason="fails on AppVeyor")
@pytest.mark.skipif(not base.HAS_SINDEX, reason="Rtree absent, skipping")
class TestFrameSindex:
def setup_method(self):
data = {
"A": range(5),
"B": range(-5, 0),
"location": [Point(x, y) for x, y in zip(range(5), range(5))],
}
self.df = GeoDataFrame(data, geometry="location")
def test_sindex(self):
self.df.crs = "epsg:4326"
assert self.df.sindex.size == 5
hits = list(self.df.sindex.intersection((2.5, 2.5, 4, 4), objects=True))
assert len(hits) == 2
assert hits[0].object == 3
def test_lazy_build(self):
assert self.df._sindex is None
assert self.df.sindex.size == 5
assert self.df._sindex is not None
def test_sindex_rebuild_on_set_geometry(self):
# First build the sindex
assert self.df.sindex is not None
self.df.set_geometry(
[Point(x, y) for x, y in zip(range(5, 10), range(5, 10))], inplace=True
)
assert self.df._sindex_generated is False
# Skip to accommodate Shapely geometries being unhashable
@pytest.mark.skip
class TestJoinSindex:
def setup_method(self):
nybb_filename = geopandas.datasets.get_path("nybb")
self.boros = read_file(nybb_filename)
def test_merge_geo(self):
# First check that we gets hits from the boros frame.
tree = self.boros.sindex
hits = tree.intersection((1012821.80, 229228.26), objects=True)
res = [self.boros.loc[hit.object]["BoroName"] for hit in hits]
assert res == ["Bronx", "Queens"]
# Check that we only get the Bronx from this view.
first = self.boros[self.boros["BoroCode"] < 3]
tree = first.sindex
hits = tree.intersection((1012821.80, 229228.26), objects=True)
res = [first.loc[hit.object]["BoroName"] for hit in hits]
assert res == ["Bronx"]
# Check that we only get Queens from this view.
second = self.boros[self.boros["BoroCode"] >= 3]
tree = second.sindex
hits = tree.intersection((1012821.80, 229228.26), objects=True)
res = ([second.loc[hit.object]["BoroName"] for hit in hits],)
assert res == ["Queens"]
# Get both the Bronx and Queens again.
merged = first.merge(second, how="outer")
assert len(merged) == 5
assert merged.sindex.size == 5
tree = merged.sindex
hits = tree.intersection((1012821.80, 229228.26), objects=True)
res = [merged.loc[hit.object]["BoroName"] for hit in hits]
assert res == ["Bronx", "Queens"]