add AccessRestrictionAffected as GeometryIndexed subclass

src/c3nav/mapdata/utils/cache/__init__.py

@@ -1,3 +1,4 @@
 from c3nav.mapdata.utils.cache.indexed import GeometryIndexed  # noqa
 from c3nav.mapdata.utils.cache.maphistory import MapHistory  # noqa
+from c3nav.mapdata.utils.cache.accessrestrictions import AccessRestrictionAffected  # noqa
 from c3nav.mapdata.utils.cache.changes import changed_geometries  # noqa

src/c3nav/mapdata/utils/cache/accessrestrictions.py

@@ -0,0 +1,104 @@
+import operator
+import struct
+from functools import reduce
+
+import numpy as np
+from shapely.geometry.base import BaseGeometry
+
+from c3nav.mapdata.utils.cache.indexed import LevelGeometryIndexed
+
+
+class AccessRestrictionAffected(LevelGeometryIndexed):
+    # metadata format:
+    # 64 times:
+    #     4 bytes (uint32): access restriction id (or 0x00 if empty)
+    # each uint64 cell contains a bitmask of restrictions.
+    # e.g.: 2^n bit set → restriction with index 2^n does apply
+    dtype = np.uint64
+    variant_id = 2
+    variant_name = 'restrictions'
+
+    def __init__(self, restrictions=None, **kwargs):
+        super().__init__(**kwargs)
+        self.restrictions = [] if restrictions is None else restrictions
+        self.restrictions_lookup = {restriction: i for i, restriction in enumerate(self.restrictions)}
+
+    @classmethod
+    def _read_metadata(cls, f, kwargs):
+        restrictions = list(struct.unpack('<'+'I'*64, f.read(4*64)))
+        while restrictions and restrictions[-1] == 0:
+            restrictions.pop()
+        kwargs['restrictions'] = restrictions
+
+    def _write_metadata(self, f):
+        f.write(struct.pack('<'+'I'*64, *self.restrictions, *((0, )*(64-len(self.restrictions)))))
+
+    @classmethod
+    def build(cls, access_restriction_affected):
+        result = cls()
+        for restriction, area in access_restriction_affected.items():
+            result[area.buffer(1)].add(restriction)
+        return result
+
+    @classmethod
+    def open(cls, filename):
+        try:
+            instance = super().open(filename)
+        except FileNotFoundError:
+            instance = cls(restrictions=[], filename=filename)
+        return instance
+
+    def _get_restriction_index(self, restriction, create=False):
+        i = self.restrictions_lookup.get(restriction)
+        if create and i is None:
+            i = len(self.restrictions)
+            self.restrictions_lookup[restriction] = i
+            self.restrictions.append(restriction)
+        return i
+
+    def __getitem__(self, selector):
+        return AccessRestrictionAffectedCells(self, selector)
+
+    def __setitem__(self, selector, value):
+        raise TypeError('__setitem__ not supported for AccessRestriction matrix')
+
+
+class AccessRestrictionAffectedCells:
+    def __init__(self, parent: AccessRestrictionAffected, selector):
+        self.parent = parent
+        self.selector = selector
+        self.values = self._get_values()
+
+    def _get_values(self):
+        return LevelGeometryIndexed.__getitem__(self.parent, self.selector)
+
+    def _set(self, values):
+        self.values = values
+        LevelGeometryIndexed.__setitem__(self.parent, self.selector, values)
+
+    def __contains__(self, restriction):
+        i = self.parent._get_restriction_index(restriction)
+        return (self.values & (2**i)).any()
+
+    def add(self, restriction):
+        if not isinstance(self.selector, BaseGeometry):
+            raise TypeError('Can only add restrictions with Geometry based selectors')
+
+        # expand array
+        bounds = self.parent._get_geometry_bounds(self.selector)
+        self.parent.fit_bounds(*bounds)
+        self._get_values()
+
+        i = self.parent._get_restriction_index(restriction, create=True)
+        self._set(self.values | (2**i))
+
+    def discard(self, restriction):
+        if not isinstance(self.selector, BaseGeometry):
+            raise TypeError('Can only discard restrictions with Geometry based selectors')
+
+        i = self.parent._get_restriction_index(restriction)
+        self._set(self.values & ((2**64-1) ^ (2**i)))
+
+    def __iter__(self):
+        all = reduce(operator.or_, self.values.tolist(), 0)
+        yield from (restriction for i, restriction in enumerate(self.parent.restrictions) if (all & 2**i))

src/c3nav/mapdata/utils/cache/indexed.py

@@ -200,6 +200,10 @@ class LevelGeometryIndexed(GeometryIndexed):
         # noinspection PyArgumentList
         return cls.open(cls.level_filename(level_id, mode), **kwargs)
 
+    def save_level(self, level_id, mode):
+        # noinspection PyArgumentList
+        return self.save(self.level_filename(level_id, mode))
+
     cached = {}
     cache_key = None
     cache_lock = threading.Lock()