Skip to content

Commit

Permalink
rtree.index.Index: avoid use of args/kwargs
Browse files Browse the repository at this point in the history
  • Loading branch information
adamjstewart committed Jan 4, 2025
1 parent 3512cac commit 28d0af1
Show file tree
Hide file tree
Showing 3 changed files with 36 additions and 73 deletions.
3 changes: 2 additions & 1 deletion benchmarks/benchmarks.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@

import rtree
from rtree import Rtree as _Rtree
from rtree.index import Property

print(f"Benchmarking Rtree-{rtree.__version__} from {Path(rtree.__file__).parent}")
print(f"Using {rtree.core.rt._name} version {rtree.core.rt.SIDX_Version().decode()}")
Expand Down Expand Up @@ -66,7 +67,7 @@ class Rtree(_Rtree):
}

index = Rtree()
disk_index = Rtree("test", overwrite=1)
disk_index = Rtree("test", properties=Property(overwrite=1))

coordinates = []
random.seed("Rtree", version=2)
Expand Down
91 changes: 26 additions & 65 deletions rtree/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import pprint
import warnings
from collections.abc import Iterator, Sequence
from typing import Any, Literal, overload
from typing import Any, Generator, Literal, overload

from . import core
from .exceptions import RTreeError
Expand Down Expand Up @@ -78,24 +78,23 @@ def _get_data(handle):
class Index:
"""An R-Tree, MVR-Tree, or TPR-Tree indexing object"""

def __init__(self, *args: Any, **kwargs: Any) -> None:
def __init__(
self,
filename: str | bytes | None = None,
stream: Any | None = None,
storage: ICustomStorage | None = None,
arrays: tuple | None = None,
interleaved: bool = True,
properties: Property | None = None,
) -> None:
"""Creates a new index
:param filename:
The first argument in the constructor is assumed to be a filename
determining that a file-based storage for the index should be used.
If the first argument is not of type basestring, it is then assumed
to be an instance of ICustomStorage or derived class.
If the first argument is neither of type basestring nor an instance
of ICustomStorage, it is then assumed to be an input index item
stream.
:param filename: a filename determining that a file-based storage for the index
should be used.
:param stream:
If the first argument in the constructor is not of type basestring,
it is assumed to be an iterable stream of data that will raise a
StopIteration. It must be in the form defined by the
:attr:`interleaved` attribute of the index. The following example
would assume :attr:`interleaved` is False::
:param stream: an iterable stream of data that will raise a StopIteration.
It must be in the form defined by the :attr:`interleaved` attribute of the
index. The following example would assume :attr:`interleaved` is False::
(id,
(minx, maxx, miny, maxy, minz, maxz, ..., ..., mink, maxk),
Expand All @@ -112,21 +111,15 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
time),
object)
:param storage:
If the first argument in the constructor is an instance of
ICustomStorage then the given custom storage is used.
:param storage: A custom storage object to be used.
:param interleaved: True or False, defaults to True.
This parameter determines the coordinate order for all methods that
:param arrays: A tuple of arrays for bulk addition.
:param interleaved: Determines the coordinate order for all methods that
take in coordinates.
:param properties: An :class:`index.Property` object.
This object sets both the creation and instantiation properties
for the object and they are passed down into libspatialindex.
A few properties are curried from instantiation parameters
for you like ``pagesize`` and ``overwrite``
to ensure compatibility with previous versions of the library. All
other properties must be set on the object.
:param properties: This object sets both the creation and instantiation
properties for the object and they are passed down into libspatialindex.
.. warning::
The coordinate ordering for all functions are sensitive the
Expand Down Expand Up @@ -194,7 +187,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
True
"""
self.properties = kwargs.get("properties", Property())
self.interleaved = interleaved
self.properties = properties or Property()

if self.properties.type == RT_TPRTree and not hasattr(
core.rt, "Index_InsertTPData"
Expand All @@ -203,46 +197,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
"TPR-Tree type not supported with version of libspatialindex"
)

# interleaved True gives 'bbox' order.
self.interleaved = bool(kwargs.get("interleaved", True))

stream = None
arrays = None
basename = None
storage = None
if args:
if isinstance(args[0], str) or isinstance(args[0], bytes):
# they sent in a filename
basename = args[0]
# they sent in a filename, stream or filename, buffers
if len(args) > 1:
if isinstance(args[1], tuple):
arrays = args[1]
else:
stream = args[1]
elif isinstance(args[0], ICustomStorage):
storage = args[0]
# they sent in a storage, stream
if len(args) > 1:
stream = args[1]
elif isinstance(args[0], tuple):
arrays = args[0]
else:
stream = args[0]

if basename:
if filename:
self.properties.storage = RT_Disk
self.properties.filename = basename
self.properties.filename = filename

# check we can read the file
f = str(basename) + "." + self.properties.idx_extension
f = str(filename) + "." + self.properties.idx_extension
p = os.path.abspath(f)

# assume if the file exists, we're not going to overwrite it
# unless the user explicitly set the property to do so
if os.path.exists(p):
self.properties.overwrite = bool(kwargs.get("overwrite", False))

# assume we're fetching the first index_id. If the user
# set it, we'll fetch that one.
if not self.properties.overwrite:
Expand All @@ -258,7 +224,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
elif storage:
self.properties.storage = RT_Custom
if storage.hasData:
self.properties.overwrite = bool(kwargs.get("overwrite", False))
if not self.properties.overwrite:
try:
self.properties.index_id
Expand All @@ -271,10 +236,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
else:
self.properties.storage = RT_Memory

ps = kwargs.get("pagesize", None)
if ps:
self.properties.pagesize = int(ps)

if stream and self.properties.type == RT_RTree:
self._exception = None
self.handle = self._create_idx_from_stream(stream)
Expand Down
15 changes: 8 additions & 7 deletions tests/test_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def stream_basic(self) -> None:
# some versions of libspatialindex screw up indexes on stream loading
# so do a very simple index check
rtree_test = rtree.index.Index(
[(1564, [0, 0, 0, 10, 10, 10], None)],
stream=[(1564, [0, 0, 0, 10, 10, 10], None)],
properties=rtree.index.Property(dimension=3),
)
assert next(rtree_test.intersection([1, 1, 1, 2, 2, 2])) == 1564
Expand Down Expand Up @@ -600,7 +600,7 @@ def test_overwrite(self) -> None:

idx = index.Index(tname)
del idx
idx = index.Index(tname, overwrite=True)
idx = index.Index(tname, properties=index.Property(overwrite=True))
assert isinstance(idx, index.Index)


Expand Down Expand Up @@ -700,7 +700,7 @@ def test_4d(self) -> None:
class IndexStream(IndexTestCase):
def test_stream_input(self) -> None:
p = index.Property()
sindex = index.Index(self.boxes15_stream(), properties=p)
sindex = index.Index(stream=self.boxes15_stream(), properties=p)
bounds = (0, 0, 60, 60)
hits = sindex.intersection(bounds)
self.assertEqual(sorted(hits), [0, 4, 16, 27, 35, 40, 47, 50, 76, 80])
Expand All @@ -726,7 +726,7 @@ def gen() -> Iterator[tuple[int, tuple[int, int, int, int], None]]:
yield (i, (1, 2, 3, 4), None)
raise TestException("raising here")

return index.Index(gen())
return index.Index(stream=gen())

self.assertRaises(TestException, create_index)

Expand Down Expand Up @@ -810,7 +810,7 @@ def test_custom_storage(self) -> None:
# we just use it here for illustrative and testing purposes.

storage = DictStorage()
r = index.Index(storage, properties=settings)
r = index.Index(storage=storage, properties=settings)

# Interestingly enough, if we take a look at the contents of our
# storage now, we can see the Rtree has already written two pages
Expand Down Expand Up @@ -849,12 +849,13 @@ def test_custom_storage_reopening(self) -> None:
settings = index.Property()
settings.writethrough = True
settings.buffering_capacity = 1
settings.overwrite = True

r1 = index.Index(storage, properties=settings, overwrite=True)
r1 = index.Index(storage=storage, properties=settings)
r1.add(555, (2, 2))
del r1
self.assertTrue(storage.hasData)

r2 = index.Index(storage, properly=settings, overwrite=False)
r2 = index.Index(storage=storage, properties=settings)
count = r2.count((0, 0, 10, 10))
self.assertEqual(count, 1)

0 comments on commit 28d0af1

Please sign in to comment.