Skip to content

Commit

Permalink
Optionally exclude input geom from results (#1166)
Browse files Browse the repository at this point in the history
* Optionally exclude input geom from results

Resolves #1106

* exclusive instead of exclude_geom
  • Loading branch information
sgillies committed Jul 15, 2021
1 parent eec996f commit a228836
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 7 deletions.
2 changes: 2 additions & 0 deletions CHANGES.txt
Expand Up @@ -10,6 +10,8 @@ Shapely 1.8 will support only Python versions >= 3.6.

New features:

- The STRtree nearest*() methods now take an optional argument that
specifies exclusion of the input geometry from results (#1115).
- A GeometryTypeError has been added to shapely.errors and is consistently
raised instead of TypeError or ValueError as in version 1.7. For backwards
compatibility, the new exception will derive from TypeError and Value error
Expand Down
34 changes: 27 additions & 7 deletions shapely/strtree.py
Expand Up @@ -21,6 +21,7 @@
import ctypes
import logging
from typing import Any, ItemsView, Iterable, Iterator, Sequence, Tuple, Union
import sys
from warnings import warn

from shapely.errors import ShapelyDeprecationWarning
Expand Down Expand Up @@ -249,7 +250,9 @@ def query(self, geom: BaseGeometry) -> Sequence[BaseGeometry]:
"""
return self.query_geoms(geom)

def nearest_item(self, geom: BaseGeometry) -> Union[Any, None]:
def nearest_item(
self, geom: BaseGeometry, exclusive: bool = False
) -> Union[Any, None]:
"""Query the tree for the node nearest to geom and get the item
stored in the node.
Expand All @@ -259,6 +262,9 @@ def nearest_item(self, geom: BaseGeometry) -> Union[Any, None]:
----------
geom : geometry object
The query geometry.
exclusive : bool, optional
Whether to exclude the item corresponding to the given geom
from results or not. Default: False.
Returns
-------
Expand Down Expand Up @@ -289,10 +295,14 @@ def nearest_item(self, geom: BaseGeometry) -> Union[Any, None]:

def callback(item1, item2, distance, userdata):
try:
callback_userdata = ctypes.cast(userdata, ctypes.py_object).value
idx = ctypes.cast(item1, ctypes.py_object).value
geom2 = ctypes.cast(item2, ctypes.py_object).value
dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double))
lgeos.GEOSDistance(self._rev[idx]._geom, geom2._geom, dist)
if callback_userdata["exclusive"] and self._rev[idx].equals(geom2):
dist[0] = sys.float_info.max
else:
lgeos.GEOSDistance(self._rev[idx]._geom, geom2._geom, dist)
return 1
except Exception:
log.exception("Caught exception")
Expand All @@ -303,19 +313,24 @@ def callback(item1, item2, distance, userdata):
ctypes.py_object(geom),
envelope._geom,
lgeos.GEOSDistanceCallback(callback),
None,
ctypes.py_object({"exclusive": exclusive}),
)
result = ctypes.cast(item, ctypes.py_object).value
return result

def nearest_geom(self, geom: BaseGeometry) -> Union[BaseGeometry, None]:
def nearest_geom(
self, geom: BaseGeometry, exclusive: bool = False
) -> Union[BaseGeometry, None]:
"""Query the tree for the node nearest to geom and get the
geometry corresponding to the item stored in the node.
Parameters
----------
geom : geometry object
The query geometry.
exclusive : bool, optional
Whether to exclude the given geom from results or not.
Default: False.
Returns
-------
Expand All @@ -325,13 +340,15 @@ def nearest_geom(self, geom: BaseGeometry) -> Union[BaseGeometry, None]:
version 2.0.
"""
item = self.nearest_item(geom)
item = self.nearest_item(geom, exclusive=exclusive)
if item is None:
return None
else:
return self._rev[item]

def nearest(self, geom: BaseGeometry) -> Union[BaseGeometry, None]:
def nearest(
self, geom: BaseGeometry, exclusive: bool = False
) -> Union[BaseGeometry, None]:
"""Query the tree for the node nearest to geom and get the
geometry corresponding to the item stored in the node.
Expand All @@ -342,6 +359,9 @@ def nearest(self, geom: BaseGeometry) -> Union[BaseGeometry, None]:
----------
geom : geometry object
The query geometry.
exclusive : bool, optional
Whether to exclude the given geom from results or not.
Default: False.
Returns
-------
Expand All @@ -351,4 +371,4 @@ def nearest(self, geom: BaseGeometry) -> Union[BaseGeometry, None]:
version 2.0.
"""
return self.nearest_geom(geom)
return self.nearest_geom(geom, exclusive=exclusive)
19 changes: 19 additions & 0 deletions tests/test_strtree.py
Expand Up @@ -177,3 +177,22 @@ def test_nearest_items(geoms, items):
with pytest.warns(ShapelyDeprecationWarning):
tree = STRtree(geoms, items)
assert tree.nearest_item(None) is None


@pytest.mark.skipif(geos_version < (3, 6, 0), reason="GEOS 3.6.0 required")
@pytest.mark.parametrize(
"geoms",
[
[
Point(0, 0.5),
Polygon([(1, 0), (2, 0), (2, 1), (1, 1)]),
Polygon([(0, 2), (1, 2), (1, 3), (0, 3)]),
]
],
)
@pytest.mark.parametrize("items", [list(range(1, 4)), list("abc")])
@pytest.mark.parametrize("query_geom", [Point(0, 0.5)])
def test_nearest_item_exclusive(geoms, items, query_geom):
with pytest.warns(ShapelyDeprecationWarning):
tree = STRtree(geoms, items)
assert tree.nearest_item(query_geom, exclusive=True) != items[0]

0 comments on commit a228836

Please sign in to comment.