import NDArray directly

This commit is contained in:
Laura Klünder 2025-04-18 01:20:50 +02:00
parent 3d2b1dff4b
commit 38c625e5ec

View file

@ -4,11 +4,8 @@ from os import PathLike
from pathlib import Path from pathlib import Path
from typing import Self, Optional, Union, TYPE_CHECKING from typing import Self, Optional, Union, TYPE_CHECKING
if TYPE_CHECKING:
from numpy.typing import NDArray
from shapely import Polygon, MultiPolygon
import numpy as np import numpy as np
from numpy.typing import NDArray
try: try:
from asgiref.local import Local as LocalContext from asgiref.local import Local as LocalContext
@ -16,6 +13,10 @@ except ImportError:
from threading import local as LocalContext from threading import local as LocalContext
if TYPE_CHECKING:
from shapely import Polygon, MultiPolygon
class GeometryIndexed: class GeometryIndexed:
# binary format (everything little-endian): # binary format (everything little-endian):
# 1 byte (uint8): variant id # 1 byte (uint8): variant id
@ -30,18 +31,18 @@ class GeometryIndexed:
variant_id = 0 variant_id = 0
def __init__(self, resolution: Optional[int] = None, x: int = 0, y: int = 0, def __init__(self, resolution: Optional[int] = None, x: int = 0, y: int = 0,
data: 'NDArray' = None, filename: str | bytes | PathLike = None): data: NDArray = None, filename: str | bytes | PathLike = None):
if resolution is None: if resolution is None:
from django.conf import settings from django.conf import settings
resolution = settings.CACHE_RESOLUTION resolution = settings.CACHE_RESOLUTION
self.resolution: int = resolution self.resolution: int = resolution
self.x = x self.x = x
self.y = y self.y = y
self.data: 'NDArray' = data if data is not None else self._get_empty_array() self.data: NDArray = data if data is not None else self._get_empty_array()
self.filename = filename self.filename = filename
@classmethod @classmethod
def _get_empty_array(cls) -> 'NDArray': def _get_empty_array(cls) -> NDArray:
return np.empty((0, 0), dtype=cls.dtype) return np.empty((0, 0), dtype=cls.dtype)
@classmethod @classmethod
@ -120,7 +121,7 @@ class GeometryIndexed:
self.y = miny self.y = miny
def get_geometry_cells(self, geometry: Union['Polygon','MultiPolygon'], def get_geometry_cells(self, geometry: Union['Polygon','MultiPolygon'],
bounds: Optional[tuple[int, int, int, int]] = None) -> 'NDArray': bounds: Optional[tuple[int, int, int, int]] = None) -> NDArray:
if bounds is None: if bounds is None:
bounds = self.get_geometry_bounds(geometry) bounds = self.get_geometry_bounds(geometry)
minx, miny, maxx, maxy = bounds minx, miny, maxx, maxy = bounds
@ -173,7 +174,7 @@ class GeometryIndexed:
raise TypeError('GeometryIndexed index must be a shapely geometry or tuple, not %s' % type(key).__name__) raise TypeError('GeometryIndexed index must be a shapely geometry or tuple, not %s' % type(key).__name__)
def __setitem__(self, key: Union['Polygon','MultiPolygon'], value: 'NDArray' | int): def __setitem__(self, key: Union['Polygon','MultiPolygon'], value: NDArray | int):
from shapely.geometry.base import BaseGeometry from shapely.geometry.base import BaseGeometry
if isinstance(key, BaseGeometry): if isinstance(key, BaseGeometry):
bounds = self.get_geometry_bounds(key) bounds = self.get_geometry_bounds(key)