"""Classes capable of reading and writing datasets Instances of these classes are called dataset objects. """ import logging from rasterio._base import get_dataset_driver, driver_can_create, driver_can_create_copy from rasterio._io import ( DatasetReaderBase, DatasetWriterBase, BufferedDatasetWriterBase, MemoryFileBase, ) from rasterio.windows import WindowMethodsMixin from rasterio.env import Env, ensure_env from rasterio.transform import TransformMethodsMixin from rasterio._path import _UnparsedPath try: from rasterio._filepath import FilePathBase except ImportError: FilePathBase = object log = logging.getLogger(__name__) class DatasetReader(DatasetReaderBase, WindowMethodsMixin, TransformMethodsMixin): """An unbuffered data and metadata reader""" def __repr__(self): return "<{} DatasetReader name='{}' mode='{}'>".format( self.closed and 'closed' or 'open', self.name, self.mode) class DatasetWriter(DatasetWriterBase, WindowMethodsMixin, TransformMethodsMixin): """An unbuffered data and metadata writer. Its methods write data directly to disk. """ def __repr__(self): return "<{} DatasetWriter name='{}' mode='{}'>".format( self.closed and 'closed' or 'open', self.name, self.mode) class BufferedDatasetWriter( BufferedDatasetWriterBase, WindowMethodsMixin, TransformMethodsMixin ): """Maintains data and metadata in a buffer, writing to disk or network only when `close()` is called. This allows incremental updates to datasets using formats that don't otherwise support updates, such as JPEG. """ def __repr__(self): return "<{} BufferedDatasetWriter name='{}' mode='{}'>".format( self.closed and 'closed' or 'open', self.name, self.mode) class MemoryFile(MemoryFileBase): """A BytesIO-like object, backed by an in-memory file. This allows formatted files to be read and written without I/O. A MemoryFile created with initial bytes becomes immutable. A MemoryFile created without initial bytes may be written to using either file-like or dataset interfaces. Examples -------- A GeoTIFF can be loaded in memory and accessed using the GeoTIFF format driver >>> with open('tests/data/RGB.byte.tif', 'rb') as f, MemoryFile(f) as memfile: ... with memfile.open() as src: ... pprint.pprint(src.profile) ... {'count': 3, 'crs': CRS({'init': 'epsg:32618'}), 'driver': 'GTiff', 'dtype': 'uint8', 'height': 718, 'interleave': 'pixel', 'nodata': 0.0, 'tiled': False, 'transform': Affine(300.0379266750948, 0.0, 101985.0, 0.0, -300.041782729805, 2826915.0), 'width': 791} """ def __init__(self, file_or_bytes=None, dirname=None, filename=None, ext=".tif"): """Create a new file in memory Parameters ---------- file_or_bytes : file-like object or bytes, optional File or bytes holding initial data. filename : str, optional An optional filename. A unique one will otherwise be generated. ext : str, optional An optional extension. Returns ------- MemoryFile """ super().__init__( file_or_bytes=file_or_bytes, dirname=dirname, filename=filename, ext=ext ) @ensure_env def open(self, driver=None, width=None, height=None, count=None, crs=None, transform=None, dtype=None, nodata=None, sharing=False, **kwargs): """Open the file and return a Rasterio dataset object. If data has already been written, the file is opened in 'r' mode. Otherwise, the file is opened in 'w' mode. Parameters ---------- Note well that there is no `path` parameter: a `MemoryFile` contains a single dataset and there is no need to specify a path. Other parameters are optional and have the same semantics as the parameters of `rasterio.open()`. """ mempath = _UnparsedPath(self.name) if self.closed: raise OSError("I/O operation on closed file.") if len(self) > 0: log.debug("VSI path: {}".format(mempath.path)) return DatasetReader(mempath, driver=driver, sharing=sharing, **kwargs) else: writer = get_writer_for_driver(driver) return writer( mempath, "w+", driver=driver, width=width, height=height, count=count, crs=crs, transform=transform, dtype=dtype, nodata=nodata, sharing=sharing, **kwargs ) def __enter__(self): return self def __exit__(self, *args): self.close() class _FilePath(FilePathBase): """A BytesIO-like object, backed by a Python file object. Examples -------- A GeoTIFF can be loaded in memory and accessed using the GeoTIFF format driver >>> with open('tests/data/RGB.byte.tif', 'rb') as f, FilePath(f) as vsi_file: ... with vsi_file.open() as src: ... pprint.pprint(src.profile) ... {'count': 3, 'crs': CRS({'init': 'epsg:32618'}), 'driver': 'GTiff', 'dtype': 'uint8', 'height': 718, 'interleave': 'pixel', 'nodata': 0.0, 'tiled': False, 'transform': Affine(300.0379266750948, 0.0, 101985.0, 0.0, -300.041782729805, 2826915.0), 'width': 791} """ def __init__(self, filelike_obj, dirname=None, filename=None): """Create a new wrapper around the provided file-like object. Parameters ---------- filelike_obj : file-like object Open file-like object. Currently only reading is supported. filename : str, optional An optional filename. A unique one will otherwise be generated. Returns ------- PythonVSIFile """ super().__init__( filelike_obj, dirname=dirname, filename=filename ) @ensure_env def open(self, driver=None, sharing=False, **kwargs): """Open the file and return a Rasterio dataset object. The provided file-like object is assumed to be readable. Writing is currently not supported. Parameters are optional and have the same semantics as the parameters of `rasterio.open()`. Returns ------- DatasetReader Raises ------ IOError If the memory file is closed. """ mempath = _UnparsedPath(self.name) if self.closed: raise IOError("I/O operation on closed file.") # Assume we were given a non-empty file-like object log.debug("VSI path: {}".format(mempath.path)) return DatasetReader(mempath, driver=driver, sharing=sharing, **kwargs) def __enter__(self): return self def __exit__(self, *args): self.close() if FilePathBase is not object: # only make this object available if the cython extension was compiled FilePath = _FilePath class ZipMemoryFile(MemoryFile): """A read-only BytesIO-like object backed by an in-memory zip file. This allows a zip file containing formatted files to be read without I/O. """ def __init__(self, file_or_bytes=None): super().__init__(file_or_bytes, ext="zip") @ensure_env def open(self, path, driver=None, sharing=False, **kwargs): """Open a dataset within the zipped stream. Parameters ---------- path : str Path to a dataset in the zip file, relative to the root of the archive. Other parameters are optional and have the same semantics as the parameters of `rasterio.open()`. Returns ------- A Rasterio dataset object """ zippath = _UnparsedPath('/vsizip{0}/{1}'.format(self.name, path.lstrip('/'))) if self.closed: raise OSError("I/O operation on closed file.") return DatasetReader(zippath, driver=driver, sharing=sharing, **kwargs) def get_writer_for_driver(driver): """Return the writer class appropriate for the specified driver.""" if not driver: raise ValueError("'driver' is required to read/write dataset.") cls = None if driver_can_create(driver): cls = DatasetWriter elif driver_can_create_copy(driver): # pragma: no branch cls = BufferedDatasetWriter return cls def get_writer_for_path(path, driver=None): """Return the writer class appropriate for the existing dataset.""" if not driver: driver = get_dataset_driver(path) return get_writer_for_driver(driver)