def get_rtree_index(
items,
objects=None,
properties=None
):
return (
RTreeIndex(generate_index_entries(items, objects=objects), properties=RTreeIndexProperty(**properties)) if properties
else RTreeIndex(generate_index_entries(items, objects=objects))
)
def __init__(self, *args, **kwargs):
self._protocol = (2 if six.sys.version_info[0] < 3 else
cpickle.HIGHEST_PROTOCOL)
idx_fp = kwargs.get('fp')
areas = kwargs.get('areas')
peril_areas = kwargs.get('peril_areas')
props = kwargs.get('properties') or copy.deepcopy(
DEFAULT_RTREE_INDEX_PROPS)
if not (idx_fp or areas or peril_areas):
self._peril_areas = self._stream = None
kwargs['properties'] = RTreeIndexProperty(**props)
super(self.__class__, self).__init__(*args, **kwargs)
elif idx_fp:
self._peril_areas = self._stream = None
_idx_fp = idx_fp
if not os.path.isabs(_idx_fp):
_idx_fp = os.path.abspath(_idx_fp)
idx_ext = props.get('idx_extension') or 'idx'
dat_ext = props.get('dat_extension') or 'dat'
if not (os.path.exists('{}.{}'.format(_idx_fp, idx_ext))
or os.path.exists('{}.{}'.format(_idx_fp, dat_ext))):
kwargs['properties'] = RTreeIndexProperty(**props)
super(self.__class__, self).__init__(_idx_fp, *args, **kwargs)
else:
self._peril_areas = OrderedDict({
pa.id: pa
for pa in
(peril_areas if peril_areas else self._get_peril_areas(areas))
})
self._stream = self._generate_index_entries(
((paid, pa.bounds)
for paid, pa in six.iteritems(self._peril_areas)),
objects=((paid, pa.bounds, pa.coordinates)
for paid, pa in six.iteritems(self._peril_areas)))
kwargs['properties'] = RTreeIndexProperty(**index_props)
super(self.__class__, self).__init__(self._stream, *args, **kwargs)
def save(self,
index_fp,
peril_areas=None,
index_props=DEFAULT_RTREE_INDEX_PROPS):
_index_fp = index_fp
if not os.path.isabs(_index_fp):
_index_fp = os.path.abspath(_index_fp)
if os.path.exists(_index_fp):
os.remove(_index_fp)
class myindex(RTreeIndex):
def __init__(self, *args, **kwargs):
self.protocol = (2 if six.sys.version_info[0] < 3 else
cpickle.HIGHEST_PROTOCOL)
super(self.__class__, self).__init__(*args, **kwargs)
def dumps(self, obj):
return cpickle.dumps(obj, protocol=self.protocol)
def loads(self, obj):
return cpickle.loads(obj)
try:
index = myindex(_index_fp,
properties=RTreeIndexProperty(**index_props))
_peril_areas = self._peril_areas or peril_areas
if not _peril_areas:
raise OasisException(
'No peril areas found in instance or in arguments - '
'this is required to write the index to file')
peril_areas_seq = None
if (isinstance(peril_areas, list)
or isinstance(peril_areas, tuple)):
peril_areas_seq = (pa for pa in peril_areas)
elif isinstance(peril_areas, types.GeneratorType):
peril_areas_seq = peril_areas
elif (isinstance(peril_areas, dict)):
peril_areas_seq = six.itervalues(peril_areas)
for pa in peril_areas_seq:
index.insert(pa.id,
pa.bounds,
obj=(pa.peril_id, pa.coverage_type, pa.id,
pa.bounds, pa.coordinates))
index.close()
except (IOError, OSError, RTreeError) as e:
raise
return _index_fp
