""" Coordinate accessors for timezonefinder. This module provides classes for accessing polygon coordinates either directly from file or from preloaded memory. """ from abc import ABC, abstractmethod import mmap from pathlib import Path from typing import Dict import numpy as np from timezonefinder import utils from timezonefinder.flatbuf.generated.polygons.PolygonCollection import ( PolygonCollection, ) from timezonefinder.flatbuf.io.polygons import ( get_polygon_collection, read_polygon_array_from_binary, ) class AbstractCoordAccessor(ABC): """Abstract base class defining the interface for coordinate accessors.""" @abstractmethod def __init__(self, coordinate_file_path: Path): """ Initialize the coordinate accessor. Args: coordinate_file_path: Path to the coordinate file """ pass @abstractmethod def __getitem__(self, idx: int) -> np.ndarray: """ Get the polygon coordinates for the given index. Args: idx: The polygon index Returns: A numpy array containing the polygon coordinates """ pass def __del__(self): """ Ensure resources are cleaned up when the object is destroyed. """ self.cleanup() @abstractmethod def cleanup(self) -> None: """Clean up resources.""" pass class FileCoordAccessor(AbstractCoordAccessor): """Accessor that reads polygon coordinates from the file on demand.""" def __init__(self, coordinate_file_path: Path): """ Initialize the file-based coordinate accessor. Args: coordinate_file_path: Path to the coordinate file """ self.coordinate_file_path = coordinate_file_path # Initialize file resources using proper resource management. try: # Use memory-mapped file for on-demand reading self.coord_file: object = open(self.coordinate_file_path, "rb") # Create memory map self.coord_buf: mmap.mmap = mmap.mmap( self.coord_file.fileno(), 0, access=mmap.ACCESS_READ ) self.polygon_collection: PolygonCollection = get_polygon_collection( self.coord_buf ) except Exception: # Clean up any partially initialized resources self.cleanup() raise def __getitem__(self, idx: int) -> np.ndarray: """ Get the polygon coordinates for the given index. Args: idx: The polygon index Returns: A numpy array containing the polygon coordinates """ return read_polygon_array_from_binary(self.polygon_collection, idx) def cleanup(self) -> None: """Clean up resources.""" # At termination utils may have been tidied up. If we're terminating we don't need to # worry about closing file handles so just avoid an exception. close_resource = getattr(utils, "close_resource", None) if close_resource is None: return # close_resource already ignores None and common close errors close_resource(getattr(self, "coord_file", None)) close_resource(getattr(self, "coord_buf", None)) # polygon_collection is just a FlatBuffers view on coord_buf and owns no resources itself class MemoryCoordAccessor(AbstractCoordAccessor): """Accessor that preloads all polygon coordinates into memory.""" def __init__(self, coordinate_file_path: Path): """ Initialize the memory-based coordinate accessor. Args: coordinate_file_path: Path to the coordinate file """ # Read entire file into memory with open(coordinate_file_path, "rb") as f: coord_buf = f.read() # Initialize polygon collection polygon_collection = get_polygon_collection(coord_buf) # Get number of polygons num_polygons = polygon_collection.PolygonsLength() # Preload all polygons self.polygons: Dict[int, np.ndarray] = {} for idx in range(num_polygons): self.polygons[idx] = read_polygon_array_from_binary(polygon_collection, idx) # Once polygons are loaded, we don't need to keep polygon_collection or coord_buf references # They'll be garbage collected def __getitem__(self, idx: int) -> np.ndarray: """ Get the polygon coordinates for the given index. Args: idx: The polygon index Returns: A numpy array containing the polygon coordinates """ return self.polygons[idx] def cleanup(self) -> None: """Clean up resources.""" del self.polygons # Just clear the dictionary, no file resources to clean up if hasattr(self, "polygons"): self.polygons.clear() def create_coord_accessor( coordinate_file_path: Path, in_memory: bool ) -> AbstractCoordAccessor: """ Factory function to create the appropriate coordinate accessor. Args: coordinate_file_path: Path to the coordinate file in_memory: Whether to use in-memory mode Returns: An instance of a coordinate accessor """ if in_memory: return MemoryCoordAccessor(coordinate_file_path) else: return FileCoordAccessor(coordinate_file_path)