rslearn.config.dataset

dataset

Classes for storing configuration of a dataset.

DType

Bases: StrEnum

Data type of a raster.

Source code in rslearn/config/dataset.py

class DType(StrEnum):
    """Data type of a raster."""

    UINT8 = "uint8"
    UINT16 = "uint16"
    UINT32 = "uint32"
    UINT64 = "uint64"
    INT8 = "int8"
    INT16 = "int16"
    INT32 = "int32"
    INT64 = "int64"
    FLOAT32 = "float32"

    def get_numpy_dtype(self) -> npt.DTypeLike:
        """Returns numpy dtype object corresponding to this DType."""
        if self == DType.UINT8:
            return np.uint8
        elif self == DType.UINT16:
            return np.uint16
        elif self == DType.UINT32:
            return np.uint32
        elif self == DType.UINT64:
            return np.uint64
        elif self == DType.INT8:
            return np.int8
        elif self == DType.INT16:
            return np.int16
        elif self == DType.INT32:
            return np.int32
        elif self == DType.INT64:
            return np.int64
        elif self == DType.FLOAT32:
            return np.float32
        raise ValueError(f"unable to handle numpy dtype {self}")

get_numpy_dtype

get_numpy_dtype() -> DTypeLike

Returns numpy dtype object corresponding to this DType.

Source code in rslearn/config/dataset.py

def get_numpy_dtype(self) -> npt.DTypeLike:
    """Returns numpy dtype object corresponding to this DType."""
    if self == DType.UINT8:
        return np.uint8
    elif self == DType.UINT16:
        return np.uint16
    elif self == DType.UINT32:
        return np.uint32
    elif self == DType.UINT64:
        return np.uint64
    elif self == DType.INT8:
        return np.int8
    elif self == DType.INT16:
        return np.int16
    elif self == DType.INT32:
        return np.int32
    elif self == DType.INT64:
        return np.int64
    elif self == DType.FLOAT32:
        return np.float32
    raise ValueError(f"unable to handle numpy dtype {self}")

ResamplingMethod

Bases: StrEnum

An enum representing the rasterio Resampling.

Source code in rslearn/config/dataset.py

class ResamplingMethod(StrEnum):
    """An enum representing the rasterio Resampling."""

    NEAREST = "nearest"
    BILINEAR = "bilinear"
    CUBIC = "cubic"
    CUBIC_SPLINE = "cubic_spline"

    def get_rasterio_resampling(self) -> Resampling:
        """Get the rasterio Resampling corresponding to this ResamplingMethod."""
        return RESAMPLING_METHODS[self]

get_rasterio_resampling

get_rasterio_resampling() -> Resampling

Get the rasterio Resampling corresponding to this ResamplingMethod.

Source code in rslearn/config/dataset.py

def get_rasterio_resampling(self) -> Resampling:
    """Get the rasterio Resampling corresponding to this ResamplingMethod."""
    return RESAMPLING_METHODS[self]

BandSetConfig

Bases: BaseModel

A configuration for a band set in a raster layer.

Each band set specifies one or more bands that should be stored together. It also specifies the storage format and dtype, the zoom offset, etc. for these bands.

Source code in rslearn/config/dataset.py

class BandSetConfig(BaseModel):
    """A configuration for a band set in a raster layer.

    Each band set specifies one or more bands that should be stored together.
    It also specifies the storage format and dtype, the zoom offset, etc. for these
    bands.
    """

    model_config = ConfigDict(extra="forbid")

    dtype: DType = Field(
        description="Pixel value type to store the data under. This is used during dataset materialize and model predict."
    )
    bands: list[str] = Field(
        default_factory=lambda: [],
        description="List of band names in this BandSetConfig. One of bands or num_bands must be set.",
    )
    num_bands: int | None = Field(
        default=None,
        description="The number of bands in this band set. The bands will be named B0, B1, B2, etc.",
    )
    format: dict[str, Any] = Field(
        default_factory=lambda: {
            "class_path": "rslearn.utils.raster_format.GeotiffRasterFormat"
        },
        description="jsonargparse configuration for the RasterFormat to store the tiles in.",
    )

    # Store images at a resolution higher or lower than the window resolution. This
    # enables keeping source data at its native resolution, either to save storage
    # space (for lower resolution data) or to retain details (for higher resolution
    # data). If positive, store data at the window resolution divided by
    # 2^(zoom_offset) (higher resolution). If negative, store data at the window
    # resolution multiplied by 2^(-zoom_offset) (lower resolution).
    zoom_offset: int = Field(
        default=0,
        description="Store data at the window resolution multiplied by 2^(-zoom_offset).",
    )

    remap: dict[str, Any] | None = Field(
        default=None,
        description="Optional configuration for a Remapper to remap pixel values.",
    )

    # Optional list of names for the different possible values of each band. The length
    # of this list must equal the number of bands. For example, [["forest", "desert"]]
    # means that it is a single-band raster where values can be 0 (forest) or 1
    # (desert).
    class_names: list[list[str]] | None = Field(
        default=None,
        description="Optional list of names for the different possible values of each band.",
    )

    # Deprecated: use nodata_value instead. Will be removed after 2026-08-01.
    nodata_vals: tuple[int | float, ...] | None = Field(
        default=None,
        description="Deprecated: per-band nodata tuple. Use nodata_value instead.",
    )

    nodata_value: int | float | None = Field(
        default=None,
        description="Optional nodata value for this band set. Used during materialization "
        "to determine which pixels are invalid when creating mosaics.",
    )

    # Optional explicit spatial dimensions for the materialized output. When set,
    # the projection resolution is adjusted so that the window's geographic extent
    # maps to exactly spatial_size pixels (height, width).
    # This is useful for coarse-resolution layers (e.g. ERA5 at 0.1 deg) where
    # only 1 pixel covers a typical window.
    spatial_size: tuple[int, int] | None = Field(
        default=None,
        description="Optional (height, width) output size. Mutually exclusive with non-zero zoom_offset.",
    )

    # Cached instantiated RasterFormat. We cache since it can take a few ms to
    # instantiate.
    _raster_format: RasterFormat | None = PrivateAttr(default=None)

    @model_validator(mode="after")
    def after_validator(self) -> "BandSetConfig":
        """Ensure the BandSetConfig is valid, and handle the num_bands field."""
        if (len(self.bands) == 0 and self.num_bands is None) or (
            len(self.bands) != 0 and self.num_bands is not None
        ):
            raise ValueError("exactly one of bands and num_bands must be specified")

        if self.num_bands is not None:
            self.bands = [f"B{band_idx}" for band_idx in range(self.num_bands)]
            self.num_bands = None

        if self.spatial_size is not None and self.zoom_offset != 0:
            raise ValueError(
                "spatial_size and non-zero zoom_offset are mutually exclusive"
            )

        if self.spatial_size is not None and any(v <= 0 for v in self.spatial_size):
            raise ValueError("spatial_size values must be positive integers")

        if self.nodata_vals is not None:
            if self.nodata_value is not None:
                raise ValueError(
                    "Cannot set both nodata_vals and nodata_value. "
                    "Use nodata_value only (nodata_vals is deprecated)."
                )
            warnings.warn(
                "nodata_vals is deprecated and will be removed after 2026-08-01. "
                "Use nodata_value (a single scalar) instead.",
                FutureWarning,
            )
            if len(self.nodata_vals) > 0:
                self.nodata_value = unique_nodata_value(self.nodata_vals)
            # Clear self.nodata_vals so it doesn't cause error in case of re-validation.
            self.nodata_vals = None

        return self

    def get_final_projection_and_bounds(
        self, projection: Projection, bounds: PixelBounds
    ) -> tuple[Projection, PixelBounds]:
        """Gets the final projection/bounds based on band set config.

        The band set config may apply a non-zero zoom offset or a fixed spatial_size
        that modifies the window's projection and bounds.

        When ``spatial_size`` is set, the projection resolution is scaled so that the
        window's geographic extent maps to exactly ``spatial_size`` pixels. The returned
        bounds will have width = spatial_size[1] and height = spatial_size[0].

        Note: the ``spatial_size`` path uses ``round()`` to compute the new pixel-space
        origin, which can shift the geographic origin by up to half of the new
        (coarser) pixel size. This is the same rounding behaviour used by
        ``ResolutionFactor.multiply_bounds`` and is negligible for coarse-resolution
        layers (e.g. ERA5 at 0.1 deg) where sub-pixel shifts are irrelevant.

        Args:
            projection: the window's projection
            bounds: the window's bounds

        Returns:
            tuple of updated projection and bounds
        """
        if self.spatial_size is not None:
            target_h, target_w = self.spatial_size
            cur_w = bounds[2] - bounds[0]
            cur_h = bounds[3] - bounds[1]

            x_factor = target_w / cur_w
            y_factor = target_h / cur_h

            new_projection = Projection(
                projection.crs,
                projection.x_resolution / x_factor,
                projection.y_resolution / y_factor,
            )
            new_bounds = (
                round(bounds[0] * x_factor),
                round(bounds[1] * y_factor),
                round(bounds[0] * x_factor) + target_w,
                round(bounds[1] * y_factor) + target_h,
            )
            return (new_projection, new_bounds)

        if self.zoom_offset >= 0:
            factor = ResolutionFactor(numerator=2**self.zoom_offset)
        else:
            factor = ResolutionFactor(denominator=2 ** (-self.zoom_offset))

        return (factor.multiply_projection(projection), factor.multiply_bounds(bounds))

    @field_validator("format", mode="before")
    @classmethod
    def convert_format_from_legacy(cls, v: dict[str, Any]) -> dict[str, Any]:
        """Support legacy format of the RasterFormat.

        The legacy format sets 'name' instead of 'class_path', and uses custom parsing
        for the init_args.
        """
        if "name" not in v:
            # New version, it is all good.
            return v

        warnings.warn(
            "`format = {'name': ...}` is deprecated; "
            "use `{'class_path': '...', 'init_args': {...}}` instead. "
            "Support will be removed after 2026-03-01.",
            FutureWarning,
        )

        legacy_name_to_class_path = {
            "image_tile": "rslearn.utils.raster_format.ImageTileRasterFormat",
            "geotiff": "rslearn.utils.raster_format.GeotiffRasterFormat",
            "single_image": "rslearn.utils.raster_format.SingleImageRasterFormat",
        }
        if v["name"] not in legacy_name_to_class_path:
            raise ValueError(
                f"could not parse legacy format with unknown raster format {v['name']}"
            )
        init_args = dict(v)
        class_path = legacy_name_to_class_path[init_args.pop("name")]

        return dict(
            class_path=class_path,
            init_args=init_args,
        )

    def instantiate_raster_format(self) -> RasterFormat:
        """Instantiate the RasterFormat specified by this BandSetConfig."""
        if self._raster_format is not None:
            return self._raster_format

        from rslearn.utils.jsonargparse import init_jsonargparse

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--raster_format", type=RasterFormat)
        cfg = parser.parse_object({"raster_format": self.format})
        self._raster_format = parser.instantiate_classes(cfg).raster_format
        return self._raster_format

after_validator

after_validator() -> BandSetConfig

Ensure the BandSetConfig is valid, and handle the num_bands field.

Source code in rslearn/config/dataset.py

@model_validator(mode="after")
def after_validator(self) -> "BandSetConfig":
    """Ensure the BandSetConfig is valid, and handle the num_bands field."""
    if (len(self.bands) == 0 and self.num_bands is None) or (
        len(self.bands) != 0 and self.num_bands is not None
    ):
        raise ValueError("exactly one of bands and num_bands must be specified")

    if self.num_bands is not None:
        self.bands = [f"B{band_idx}" for band_idx in range(self.num_bands)]
        self.num_bands = None

    if self.spatial_size is not None and self.zoom_offset != 0:
        raise ValueError(
            "spatial_size and non-zero zoom_offset are mutually exclusive"
        )

    if self.spatial_size is not None and any(v <= 0 for v in self.spatial_size):
        raise ValueError("spatial_size values must be positive integers")

    if self.nodata_vals is not None:
        if self.nodata_value is not None:
            raise ValueError(
                "Cannot set both nodata_vals and nodata_value. "
                "Use nodata_value only (nodata_vals is deprecated)."
            )
        warnings.warn(
            "nodata_vals is deprecated and will be removed after 2026-08-01. "
            "Use nodata_value (a single scalar) instead.",
            FutureWarning,
        )
        if len(self.nodata_vals) > 0:
            self.nodata_value = unique_nodata_value(self.nodata_vals)
        # Clear self.nodata_vals so it doesn't cause error in case of re-validation.
        self.nodata_vals = None

    return self

get_final_projection_and_bounds

get_final_projection_and_bounds(projection: Projection, bounds: PixelBounds) -> tuple[Projection, PixelBounds]

Gets the final projection/bounds based on band set config.

The band set config may apply a non-zero zoom offset or a fixed spatial_size that modifies the window's projection and bounds.

When spatial_size is set, the projection resolution is scaled so that the window's geographic extent maps to exactly spatial_size pixels. The returned bounds will have width = spatial_size[1] and height = spatial_size[0].

Note: the spatial_size path uses round() to compute the new pixel-space origin, which can shift the geographic origin by up to half of the new (coarser) pixel size. This is the same rounding behaviour used by ResolutionFactor.multiply_bounds and is negligible for coarse-resolution layers (e.g. ERA5 at 0.1 deg) where sub-pixel shifts are irrelevant.

Parameters:

Name	Type	Description	Default
projection	Projection	the window's projection	required
bounds	PixelBounds	the window's bounds	required

Returns:

Type	Description
tuple[Projection, PixelBounds]	tuple of updated projection and bounds

Source code in rslearn/config/dataset.py

def get_final_projection_and_bounds(
    self, projection: Projection, bounds: PixelBounds
) -> tuple[Projection, PixelBounds]:
    """Gets the final projection/bounds based on band set config.

    The band set config may apply a non-zero zoom offset or a fixed spatial_size
    that modifies the window's projection and bounds.

    When ``spatial_size`` is set, the projection resolution is scaled so that the
    window's geographic extent maps to exactly ``spatial_size`` pixels. The returned
    bounds will have width = spatial_size[1] and height = spatial_size[0].

    Note: the ``spatial_size`` path uses ``round()`` to compute the new pixel-space
    origin, which can shift the geographic origin by up to half of the new
    (coarser) pixel size. This is the same rounding behaviour used by
    ``ResolutionFactor.multiply_bounds`` and is negligible for coarse-resolution
    layers (e.g. ERA5 at 0.1 deg) where sub-pixel shifts are irrelevant.

    Args:
        projection: the window's projection
        bounds: the window's bounds

    Returns:
        tuple of updated projection and bounds
    """
    if self.spatial_size is not None:
        target_h, target_w = self.spatial_size
        cur_w = bounds[2] - bounds[0]
        cur_h = bounds[3] - bounds[1]

        x_factor = target_w / cur_w
        y_factor = target_h / cur_h

        new_projection = Projection(
            projection.crs,
            projection.x_resolution / x_factor,
            projection.y_resolution / y_factor,
        )
        new_bounds = (
            round(bounds[0] * x_factor),
            round(bounds[1] * y_factor),
            round(bounds[0] * x_factor) + target_w,
            round(bounds[1] * y_factor) + target_h,
        )
        return (new_projection, new_bounds)

    if self.zoom_offset >= 0:
        factor = ResolutionFactor(numerator=2**self.zoom_offset)
    else:
        factor = ResolutionFactor(denominator=2 ** (-self.zoom_offset))

    return (factor.multiply_projection(projection), factor.multiply_bounds(bounds))

convert_format_from_legacy classmethod

convert_format_from_legacy(v: dict[str, Any]) -> dict[str, Any]

Support legacy format of the RasterFormat.

The legacy format sets 'name' instead of 'class_path', and uses custom parsing for the init_args.

Source code in rslearn/config/dataset.py

@field_validator("format", mode="before")
@classmethod
def convert_format_from_legacy(cls, v: dict[str, Any]) -> dict[str, Any]:
    """Support legacy format of the RasterFormat.

    The legacy format sets 'name' instead of 'class_path', and uses custom parsing
    for the init_args.
    """
    if "name" not in v:
        # New version, it is all good.
        return v

    warnings.warn(
        "`format = {'name': ...}` is deprecated; "
        "use `{'class_path': '...', 'init_args': {...}}` instead. "
        "Support will be removed after 2026-03-01.",
        FutureWarning,
    )

    legacy_name_to_class_path = {
        "image_tile": "rslearn.utils.raster_format.ImageTileRasterFormat",
        "geotiff": "rslearn.utils.raster_format.GeotiffRasterFormat",
        "single_image": "rslearn.utils.raster_format.SingleImageRasterFormat",
    }
    if v["name"] not in legacy_name_to_class_path:
        raise ValueError(
            f"could not parse legacy format with unknown raster format {v['name']}"
        )
    init_args = dict(v)
    class_path = legacy_name_to_class_path[init_args.pop("name")]

    return dict(
        class_path=class_path,
        init_args=init_args,
    )

instantiate_raster_format

instantiate_raster_format() -> RasterFormat

Instantiate the RasterFormat specified by this BandSetConfig.

Source code in rslearn/config/dataset.py

def instantiate_raster_format(self) -> RasterFormat:
    """Instantiate the RasterFormat specified by this BandSetConfig."""
    if self._raster_format is not None:
        return self._raster_format

    from rslearn.utils.jsonargparse import init_jsonargparse

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--raster_format", type=RasterFormat)
    cfg = parser.parse_object({"raster_format": self.format})
    self._raster_format = parser.instantiate_classes(cfg).raster_format
    return self._raster_format

SpaceMode

Bases: StrEnum

Spatial matching mode when looking up items corresponding to a window.

Source code in rslearn/config/dataset.py

class SpaceMode(StrEnum):
    """Spatial matching mode when looking up items corresponding to a window."""

    CONTAINS = "CONTAINS"
    """Items must contain the entire window."""

    INTERSECTS = "INTERSECTS"
    """Items must overlap any portion of the window."""

    MOSAIC = "MOSAIC"
    """Groups of items should be computed that cover the entire window.

    During materialization, items in each group are merged to form a mosaic in the
    dataset.
    """

    PER_PERIOD_MOSAIC = "PER_PERIOD_MOSAIC"
    """Deprecated: use MOSAIC with period_duration instead. Will be removed after 2026-05-01."""

    SINGLE_COMPOSITE = "SINGLE_COMPOSITE"
    """Put all intersecting items into a single group.

    This can be used together with compositing method to create one composite for the
    layer.
    """

CONTAINS class-attribute instance-attribute

CONTAINS = 'CONTAINS'

Items must contain the entire window.

INTERSECTS class-attribute instance-attribute

INTERSECTS = 'INTERSECTS'

Items must overlap any portion of the window.

MOSAIC class-attribute instance-attribute

MOSAIC = 'MOSAIC'

Groups of items should be computed that cover the entire window.

During materialization, items in each group are merged to form a mosaic in the dataset.

PER_PERIOD_MOSAIC class-attribute instance-attribute

PER_PERIOD_MOSAIC = 'PER_PERIOD_MOSAIC'

Deprecated: use MOSAIC with period_duration instead. Will be removed after 2026-05-01.

SINGLE_COMPOSITE class-attribute instance-attribute

SINGLE_COMPOSITE = 'SINGLE_COMPOSITE'

Put all intersecting items into a single group.

This can be used together with compositing method to create one composite for the layer.

TimeMode

Bases: StrEnum

Temporal matching mode when looking up items corresponding to a window.

Note: setting this is deprecated. It should always be set to TimeMode.WITHIN, and TimeMode will be removed in a future release.

Source code in rslearn/config/dataset.py

class TimeMode(StrEnum):
    """Temporal matching mode when looking up items corresponding to a window.

    Note: setting this is deprecated. It should always be set to TimeMode.WITHIN, and
    TimeMode will be removed in a future release.
    """

    WITHIN = "WITHIN"
    """Items must be within the window time range."""

    NEAREST = "NEAREST"
    """Deprecated: Select items closest to the window time range, up to max_matches."""

    BEFORE = "BEFORE"
    """Deprecated: Select items before the end of the window time range, up to max_matches."""

    AFTER = "AFTER"
    """Deprecated: Select items after the start of the window time range, up to max_matches."""

WITHIN class-attribute instance-attribute

WITHIN = 'WITHIN'

Items must be within the window time range.

NEAREST class-attribute instance-attribute

NEAREST = 'NEAREST'

Deprecated: Select items closest to the window time range, up to max_matches.

BEFORE class-attribute instance-attribute

BEFORE = 'BEFORE'

Deprecated: Select items before the end of the window time range, up to max_matches.

AFTER class-attribute instance-attribute

AFTER = 'AFTER'

Deprecated: Select items after the start of the window time range, up to max_matches.

QueryConfig

Bases: BaseModel

A configuration for querying items in a data source.

Source code in rslearn/config/dataset.py

class QueryConfig(BaseModel):
    """A configuration for querying items in a data source."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    space_mode: SpaceMode = Field(
        default=SpaceMode.MOSAIC,
        description="Specifies how items should be matched with windows spatially.",
    )
    time_mode: TimeMode = Field(
        default=TimeMode.WITHIN,
        description="Deprecated: This option will be removed. Only WITHIN is supported.",
        exclude=True,
    )

    @model_validator(mode="after")
    def _warn_deprecated_fields(self) -> "QueryConfig":
        if "time_mode" in self.model_fields_set:
            warnings.warn(
                "time_mode is deprecated and will be removed in a future version. "
                "Remove it from your config (WITHIN is the only supported behavior).",
                FutureWarning,
            )
        if self.space_mode == SpaceMode.PER_PERIOD_MOSAIC:
            warnings.warn(
                "SpaceMode.PER_PERIOD_MOSAIC is deprecated and will be removed after "
                "2026-05-01. Use SpaceMode.MOSAIC with period_duration instead.",
                FutureWarning,
            )
        return self

    # Minimum number of item groups. If there are fewer than this many matches, then no
    # matches will be returned. This can be used to prevent unnecessary data ingestion
    # if the user plans to discard windows that do not have a sufficient amount of data.
    min_matches: int = Field(
        default=0, description="The minimum number of item groups."
    )

    max_matches: int = Field(
        default=1,
        description="The maximum number of item groups. When period_duration is set, this "
        "controls the maximum number of sub-periods for which item groups are created,"
        "and within each sub-period, the matching strategy is applied with an effective "
        "max_matches of 1 (i.e. one item group per sub-period).",
    )
    period_duration: Annotated[
        timedelta | None,
        BeforeValidator(ensure_optional_timedelta),
        PlainSerializer(serialize_optional_timedelta),
    ] = Field(
        default=None,
        description="If set, split the window's time range into sub-periods of this duration. "
        "The sub-period splitting takes effect before the SpaceMode, i.e., SpaceMode matching "
        "operates within the sub-periods. Each sub-period produces a single separate item group, "
        "up to max_matches total groups/periods.",
    )
    mosaic_compositing_overlaps: int = Field(
        default=1,
        description="For MOSAIC and PER_PERIOD_MOSAIC modes, the number of overlapping items "
        "wanted within each item group covering the window. Set to 1 for a single coverage "
        "(default mosaic behavior), or higher for compositing multiple overlapping items."
        "with mean or median compositing method.",
    )
    per_period_mosaic_reverse_time_order: bool = Field(
        default=True,
        description="For PER_PERIOD_MOSAIC mode, whether to return item groups in reverse "
        "temporal order (most recent first). Set to False for chronological order (oldest first). "
        "Default True is deprecated and will change to False with error if still unset or set True "
        "after 2026-04-01.",
    )

DataSourceConfig

Bases: BaseModel

Configuration for a DataSource in a dataset layer.

Source code in rslearn/config/dataset.py

class DataSourceConfig(BaseModel):
    """Configuration for a DataSource in a dataset layer."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    class_path: str = Field(description="Class path for the data source.")
    init_args: dict[str, Any] = Field(
        default_factory=lambda: {},
        description="jsonargparse init args for the data source.",
    )
    query_config: QueryConfig = Field(
        default_factory=lambda: QueryConfig(),
        description="QueryConfig specifying how to match items with windows.",
    )
    time_offset: Annotated[
        timedelta | None,
        BeforeValidator(ensure_optional_timedelta),
        PlainSerializer(serialize_optional_timedelta),
    ] = Field(
        default=None,
        description="Optional timedelta to add to the window's time range before matching.",
    )
    duration: Annotated[
        timedelta | None,
        BeforeValidator(ensure_optional_timedelta),
        PlainSerializer(serialize_optional_timedelta),
    ] = Field(
        default=None,
        description="Optional, if the window's time range is (t0, t1), then update to (t0, t0 + duration).",
    )
    ingest: bool = Field(
        default=True,
        description="Whether to ingest this layer (default True). If False, it will be directly materialized without ingestion.",
    )

    def get_request_time_range(
        self, window_time_range: tuple[datetime, datetime] | None
    ) -> tuple[datetime, datetime] | None:
        """Apply time_offset and duration to a window time range.

        This converts a window's time range into the request time range that should be
        used during prepare and materialize.

        Args:
            window_time_range: the window's original time range, or None.

        Returns:
            The adjusted time range, or None if the input was None.
        """
        if window_time_range is None:
            return None
        result = window_time_range
        if self.time_offset:
            result = (result[0] + self.time_offset, result[1] + self.time_offset)
        if self.duration:
            result = (result[0], result[0] + self.duration)
        return result

    @model_validator(mode="before")
    @classmethod
    def convert_from_legacy(cls, d: dict[str, Any]) -> dict[str, Any]:
        """Support legacy format of the DataSourceConfig.

        The legacy format sets 'name' instead of 'class_path', and mixes the arguments
        for the DataSource in with the DataSourceConfig keys.
        """
        if "name" not in d:
            # New version, it is all good.
            return d

        warnings.warn(
            "`Data source configuration {'name': ...}` is deprecated; "
            "use `{'class_path': '...', 'init_args': {...}, ...}` instead. "
            "Support will be removed after 2026-03-01.",
            FutureWarning,
        )

        # Split the dict into the base config that is in the pydantic model, and the
        # source-specific options that should be moved to init_args dict.
        class_path = d["name"]
        base_config: dict[str, Any] = {}
        ds_init_args: dict[str, Any] = {}
        for k, v in d.items():
            if k == "name":
                continue
            if k in cls.model_fields:
                base_config[k] = v
            else:
                ds_init_args[k] = v

        # Some legacy configs erroneously specify these keys, which are now caught by
        # validation. But we still want those specific legacy configs to work.
        if (
            class_path == "rslearn.data_sources.planetary_computer.Sentinel2"
            and "max_cloud_cover" in ds_init_args
        ):
            warnings.warn(
                "Data source configuration specifies invalid 'max_cloud_cover' option."
                "Support for ignoring this option will be removed after 2026-03-01.",
                FutureWarning,
            )
            del ds_init_args["max_cloud_cover"]

        base_config["class_path"] = class_path
        base_config["init_args"] = ds_init_args
        return base_config

get_request_time_range

get_request_time_range(window_time_range: tuple[datetime, datetime] | None) -> tuple[datetime, datetime] | None

Apply time_offset and duration to a window time range.

This converts a window's time range into the request time range that should be used during prepare and materialize.

Parameters:

Name	Type	Description	Default
window_time_range	tuple[datetime, datetime] | None	the window's original time range, or None.	required

Returns:

Type	Description
tuple[datetime, datetime] | None	The adjusted time range, or None if the input was None.

Source code in rslearn/config/dataset.py

def get_request_time_range(
    self, window_time_range: tuple[datetime, datetime] | None
) -> tuple[datetime, datetime] | None:
    """Apply time_offset and duration to a window time range.

    This converts a window's time range into the request time range that should be
    used during prepare and materialize.

    Args:
        window_time_range: the window's original time range, or None.

    Returns:
        The adjusted time range, or None if the input was None.
    """
    if window_time_range is None:
        return None
    result = window_time_range
    if self.time_offset:
        result = (result[0] + self.time_offset, result[1] + self.time_offset)
    if self.duration:
        result = (result[0], result[0] + self.duration)
    return result

convert_from_legacy classmethod

convert_from_legacy(d: dict[str, Any]) -> dict[str, Any]

Support legacy format of the DataSourceConfig.

The legacy format sets 'name' instead of 'class_path', and mixes the arguments for the DataSource in with the DataSourceConfig keys.

Source code in rslearn/config/dataset.py

@model_validator(mode="before")
@classmethod
def convert_from_legacy(cls, d: dict[str, Any]) -> dict[str, Any]:
    """Support legacy format of the DataSourceConfig.

    The legacy format sets 'name' instead of 'class_path', and mixes the arguments
    for the DataSource in with the DataSourceConfig keys.
    """
    if "name" not in d:
        # New version, it is all good.
        return d

    warnings.warn(
        "`Data source configuration {'name': ...}` is deprecated; "
        "use `{'class_path': '...', 'init_args': {...}, ...}` instead. "
        "Support will be removed after 2026-03-01.",
        FutureWarning,
    )

    # Split the dict into the base config that is in the pydantic model, and the
    # source-specific options that should be moved to init_args dict.
    class_path = d["name"]
    base_config: dict[str, Any] = {}
    ds_init_args: dict[str, Any] = {}
    for k, v in d.items():
        if k == "name":
            continue
        if k in cls.model_fields:
            base_config[k] = v
        else:
            ds_init_args[k] = v

    # Some legacy configs erroneously specify these keys, which are now caught by
    # validation. But we still want those specific legacy configs to work.
    if (
        class_path == "rslearn.data_sources.planetary_computer.Sentinel2"
        and "max_cloud_cover" in ds_init_args
    ):
        warnings.warn(
            "Data source configuration specifies invalid 'max_cloud_cover' option."
            "Support for ignoring this option will be removed after 2026-03-01.",
            FutureWarning,
        )
        del ds_init_args["max_cloud_cover"]

    base_config["class_path"] = class_path
    base_config["init_args"] = ds_init_args
    return base_config

LayerType

Bases: StrEnum

The layer type (raster or vector).

Source code in rslearn/config/dataset.py

class LayerType(StrEnum):
    """The layer type (raster or vector)."""

    RASTER = "raster"
    VECTOR = "vector"

CompositingMethod

Bases: StrEnum

Method how to select pixels for the composite from corresponding items of a window.

For MEAN and MEDIAN modes, mosaic_compositing_overlaps (in the QueryConfig) should be set higher than 1 so that rslearn creates item groups during prepare that cover the window with multiple overlaps. At each pixel/band, the mean and median can then be computed across items in each group that cover that pixel.

Source code in rslearn/config/dataset.py

class CompositingMethod(StrEnum):
    """Method how to select pixels for the composite from corresponding items of a window.

    For MEAN and MEDIAN modes, mosaic_compositing_overlaps (in the QueryConfig) should
    be set higher than 1 so that rslearn creates item groups during prepare that cover
    the window with multiple overlaps. At each pixel/band, the mean and median can then
    be computed across items in each group that cover that pixel.
    """

    FIRST_VALID = "FIRST_VALID"
    """Select first valid pixel in order of corresponding items (might be sorted)"""

    MEAN = "MEAN"
    """Select per-pixel mean value of corresponding items of a window"""

    MEDIAN = "MEDIAN"
    """Select per-pixel median value of corresponding items of a window"""

    SPATIAL_MOSAIC_TEMPORAL_STACK = "SPATIAL_MOSAIC_TEMPORAL_STACK"
    """Spatial first-valid compositing per timestep, stacked along T.

    Items, which can contain multi-temporal rasters, are spatially composited using
    first-valid logic within each timestep, but timesteps are stacked. The result is a
    (C, T, H, W) RasterArray whose T dimension spans the union of all item timesteps,
    clipped to the window time range.
    """

    TEMPORAL_MEAN = "TEMPORAL_MEAN"
    """Reduce a multi-temporal raster stack to one timestep via temporal mean."""

    TEMPORAL_MAX = "TEMPORAL_MAX"
    """Reduce a multi-temporal raster stack to one timestep via temporal max."""

    TEMPORAL_MIN = "TEMPORAL_MIN"
    """Reduce a multi-temporal raster stack to one timestep via temporal min."""

FIRST_VALID class-attribute instance-attribute

FIRST_VALID = 'FIRST_VALID'

Select first valid pixel in order of corresponding items (might be sorted)

MEAN class-attribute instance-attribute

MEAN = 'MEAN'

Select per-pixel mean value of corresponding items of a window

MEDIAN class-attribute instance-attribute

MEDIAN = 'MEDIAN'

Select per-pixel median value of corresponding items of a window

SPATIAL_MOSAIC_TEMPORAL_STACK class-attribute instance-attribute

SPATIAL_MOSAIC_TEMPORAL_STACK = 'SPATIAL_MOSAIC_TEMPORAL_STACK'

Spatial first-valid compositing per timestep, stacked along T.

Items, which can contain multi-temporal rasters, are spatially composited using first-valid logic within each timestep, but timesteps are stacked. The result is a (C, T, H, W) RasterArray whose T dimension spans the union of all item timesteps, clipped to the window time range.

TEMPORAL_MEAN class-attribute instance-attribute

TEMPORAL_MEAN = 'TEMPORAL_MEAN'

Reduce a multi-temporal raster stack to one timestep via temporal mean.

TEMPORAL_MAX class-attribute instance-attribute

TEMPORAL_MAX = 'TEMPORAL_MAX'

Reduce a multi-temporal raster stack to one timestep via temporal max.

TEMPORAL_MIN class-attribute instance-attribute

TEMPORAL_MIN = 'TEMPORAL_MIN'

Reduce a multi-temporal raster stack to one timestep via temporal min.

LayerConfig

Bases: BaseModel

Configuration of a layer in a dataset.

Source code in rslearn/config/dataset.py

class LayerConfig(BaseModel):
    """Configuration of a layer in a dataset."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    type: LayerType = Field(description="The LayerType (raster or vector).")
    data_source: DataSourceConfig | None = Field(
        default=None,
        description="Optional DataSourceConfig if this layer is retrievable.",
    )
    alias: str | None = Field(
        default=None, description="Alias for this layer to use in the tile store."
    )

    # Raster layer options.
    band_sets: list[BandSetConfig] = Field(
        default_factory=lambda: [],
        description="For raster layers, the bands to store in this layer.",
    )
    resampling_method: ResamplingMethod = Field(
        default=ResamplingMethod.BILINEAR,
        description="For raster layers, how to resample rasters (if neeed), default bilinear resampling.",
    )
    compositing_method: CompositingMethod | dict[str, Any] = Field(
        default=CompositingMethod.FIRST_VALID,
        description=(
            "For raster layers, how to compute pixel values in the composite of each "
            "window's items. Either a built-in enum name (e.g. FIRST_VALID) or a "
            "jsonargparse class_path/init_args dict for a custom Compositor subclass."
        ),
    )

    # Vector layer options.
    vector_format: dict[str, Any] = Field(
        default_factory=lambda: {
            "class_path": "rslearn.utils.vector_format.GeojsonVectorFormat"
        },
        description="For vector layers, the jsonargparse configuration for the VectorFormat.",
    )
    class_property_name: str | None = Field(
        default=None,
        description="Optional metadata field indicating that the GeoJSON features contain a property that corresponds to a class label, and this is the name of that property.",
    )
    class_names: list[str] | None = Field(
        default=None,
        description="The list of classes that the class_property_name property could be set to.",
    )

    @model_validator(mode="after")
    def after_validator(self) -> "LayerConfig":
        """Ensure the LayerConfig is valid."""
        if self.type == LayerType.RASTER and len(self.band_sets) == 0:
            raise ValueError(
                "band sets must be specified and non-empty for raster layers"
            )

        return self

    def __hash__(self) -> int:
        """Return a hash of this LayerConfig."""
        return hash(json.dumps(self.model_dump(mode="json"), sort_keys=True))

    def __eq__(self, other: Any) -> bool:
        """Returns whether other is the same as this LayerConfig.

        Args:
            other: the other object to compare.
        """
        if not isinstance(other, LayerConfig):
            return False
        return self.model_dump() == other.model_dump()

    @functools.cache
    def instantiate_data_source(self, ds_path: UPath | None = None) -> "DataSource":
        """Instantiate the data source specified by this config.

        Args:
            ds_path: optional dataset path to include in the DataSourceContext.

        Returns:
            the DataSource object.
        """
        from rslearn.data_sources.data_source import DataSource, DataSourceContext
        from rslearn.utils.jsonargparse import data_source_context_serializer

        logger.debug("getting a data source for dataset at %s", ds_path)
        if self.data_source is None:
            raise ValueError("This layer does not specify a data source")

        # Inject the DataSourceContext into the args.
        context = DataSourceContext(
            ds_path=ds_path,
            layer_config=self,
        )
        ds_config: dict[str, Any] = {
            "class_path": self.data_source.class_path,
            "init_args": copy.deepcopy(self.data_source.init_args),
        }
        ds_config["init_args"]["context"] = data_source_context_serializer(context)

        # Now we can parse with jsonargparse.
        from rslearn.utils.jsonargparse import (
            data_source_context_serializer,
            init_jsonargparse,
        )

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--data_source", type=DataSource)
        cfg = parser.parse_object({"data_source": ds_config})
        data_source = parser.instantiate_classes(cfg).data_source
        return data_source

    def instantiate_vector_format(self) -> VectorFormat:
        """Instantiate the vector format specified by this config."""
        if self.type != LayerType.VECTOR:
            raise ValueError(
                f"cannot instantiate vector format for layer with type {self.type}"
            )

        from rslearn.utils.jsonargparse import init_jsonargparse

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--vector_format", type=VectorFormat)
        cfg = parser.parse_object({"vector_format": self.vector_format})
        vector_format = parser.instantiate_classes(cfg).vector_format
        return vector_format

    def instantiate_compositor(self) -> "Compositor":
        """Instantiate the Compositor for this layer's compositing_method.

        Returns:
            A Compositor instance -- either a built-in one resolved from the
            CompositingMethod enum, or a custom one instantiated via jsonargparse
            from a class_path/init_args dict.
        """
        if isinstance(self.compositing_method, CompositingMethod):
            from rslearn.dataset.compositing import BUILTIN_COMPOSITORS

            return BUILTIN_COMPOSITORS[self.compositing_method]

        from rslearn.dataset.compositing import Compositor
        from rslearn.utils.jsonargparse import init_jsonargparse

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--compositor", type=Compositor)
        cfg = parser.parse_object({"compositor": self.compositing_method})
        return parser.instantiate_classes(cfg).compositor

after_validator

after_validator() -> LayerConfig

Ensure the LayerConfig is valid.

Source code in rslearn/config/dataset.py

@model_validator(mode="after")
def after_validator(self) -> "LayerConfig":
    """Ensure the LayerConfig is valid."""
    if self.type == LayerType.RASTER and len(self.band_sets) == 0:
        raise ValueError(
            "band sets must be specified and non-empty for raster layers"
        )

    return self

instantiate_data_source cached

instantiate_data_source(ds_path: UPath | None = None) -> DataSource

Instantiate the data source specified by this config.

Parameters:

Name	Type	Description	Default
ds_path	UPath | None	optional dataset path to include in the DataSourceContext.	None

Returns:

Type	Description
DataSource	the DataSource object.

Source code in rslearn/config/dataset.py

@functools.cache
def instantiate_data_source(self, ds_path: UPath | None = None) -> "DataSource":
    """Instantiate the data source specified by this config.

    Args:
        ds_path: optional dataset path to include in the DataSourceContext.

    Returns:
        the DataSource object.
    """
    from rslearn.data_sources.data_source import DataSource, DataSourceContext
    from rslearn.utils.jsonargparse import data_source_context_serializer

    logger.debug("getting a data source for dataset at %s", ds_path)
    if self.data_source is None:
        raise ValueError("This layer does not specify a data source")

    # Inject the DataSourceContext into the args.
    context = DataSourceContext(
        ds_path=ds_path,
        layer_config=self,
    )
    ds_config: dict[str, Any] = {
        "class_path": self.data_source.class_path,
        "init_args": copy.deepcopy(self.data_source.init_args),
    }
    ds_config["init_args"]["context"] = data_source_context_serializer(context)

    # Now we can parse with jsonargparse.
    from rslearn.utils.jsonargparse import (
        data_source_context_serializer,
        init_jsonargparse,
    )

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--data_source", type=DataSource)
    cfg = parser.parse_object({"data_source": ds_config})
    data_source = parser.instantiate_classes(cfg).data_source
    return data_source

instantiate_vector_format

instantiate_vector_format() -> VectorFormat

Instantiate the vector format specified by this config.

Source code in rslearn/config/dataset.py

def instantiate_vector_format(self) -> VectorFormat:
    """Instantiate the vector format specified by this config."""
    if self.type != LayerType.VECTOR:
        raise ValueError(
            f"cannot instantiate vector format for layer with type {self.type}"
        )

    from rslearn.utils.jsonargparse import init_jsonargparse

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--vector_format", type=VectorFormat)
    cfg = parser.parse_object({"vector_format": self.vector_format})
    vector_format = parser.instantiate_classes(cfg).vector_format
    return vector_format

instantiate_compositor

instantiate_compositor() -> Compositor

Instantiate the Compositor for this layer's compositing_method.

Returns:

Type	Description
Compositor	A Compositor instance -- either a built-in one resolved from the
Compositor	CompositingMethod enum, or a custom one instantiated via jsonargparse
Compositor	from a class_path/init_args dict.

Source code in rslearn/config/dataset.py

def instantiate_compositor(self) -> "Compositor":
    """Instantiate the Compositor for this layer's compositing_method.

    Returns:
        A Compositor instance -- either a built-in one resolved from the
        CompositingMethod enum, or a custom one instantiated via jsonargparse
        from a class_path/init_args dict.
    """
    if isinstance(self.compositing_method, CompositingMethod):
        from rslearn.dataset.compositing import BUILTIN_COMPOSITORS

        return BUILTIN_COMPOSITORS[self.compositing_method]

    from rslearn.dataset.compositing import Compositor
    from rslearn.utils.jsonargparse import init_jsonargparse

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--compositor", type=Compositor)
    cfg = parser.parse_object({"compositor": self.compositing_method})
    return parser.instantiate_classes(cfg).compositor

StorageConfig

Bases: BaseModel

Configuration for the WindowStorageFactory (window metadata storage backend).

Source code in rslearn/config/dataset.py

class StorageConfig(BaseModel):
    """Configuration for the WindowStorageFactory (window metadata storage backend)."""

    model_config = ConfigDict(frozen=True, extra="forbid")

    class_path: str = Field(
        default="rslearn.dataset.storage.file.FileWindowStorageFactory",
        description="Class path for the WindowStorageFactory.",
    )
    init_args: dict[str, Any] = Field(
        default_factory=lambda: {},
        description="jsonargparse init args for the WindowStorageFactory.",
    )

    def instantiate_window_storage_factory(self) -> "WindowStorageFactory":
        """Instantiate the WindowStorageFactory specified by this config."""
        from rslearn.dataset.storage.storage import WindowStorageFactory
        from rslearn.utils.jsonargparse import init_jsonargparse

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--wsf", type=WindowStorageFactory)
        cfg = parser.parse_object(
            {
                "wsf": dict(
                    class_path=self.class_path,
                    init_args=self.init_args,
                )
            }
        )
        wsf = parser.instantiate_classes(cfg).wsf
        return wsf

instantiate_window_storage_factory

instantiate_window_storage_factory() -> WindowStorageFactory

Instantiate the WindowStorageFactory specified by this config.

Source code in rslearn/config/dataset.py

def instantiate_window_storage_factory(self) -> "WindowStorageFactory":
    """Instantiate the WindowStorageFactory specified by this config."""
    from rslearn.dataset.storage.storage import WindowStorageFactory
    from rslearn.utils.jsonargparse import init_jsonargparse

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--wsf", type=WindowStorageFactory)
    cfg = parser.parse_object(
        {
            "wsf": dict(
                class_path=self.class_path,
                init_args=self.init_args,
            )
        }
    )
    wsf = parser.instantiate_classes(cfg).wsf
    return wsf

WindowDataStorageConfig

Bases: BaseModel

Configuration for the WindowDataStorage of the dataset.

The WindowDataStorage controls the on-disk layout of materialized raster and vector data within each window.

Source code in rslearn/config/dataset.py

class WindowDataStorageConfig(BaseModel):
    """Configuration for the WindowDataStorage of the dataset.

    The WindowDataStorage controls the on-disk layout of materialized raster
    and vector data within each window.
    """

    model_config = ConfigDict(frozen=True, extra="forbid")

    class_path: str = Field(
        default="rslearn.dataset.window_data_storage.per_item_group.PerItemGroupStorageFactory",
        description="Class path for the WindowDataStorageFactory.",
    )
    init_args: dict[str, Any] = Field(
        default_factory=lambda: {},
        description="jsonargparse init args for the WindowDataStorageFactory.",
    )

    def instantiate_factory(self) -> "WindowDataStorageFactory":
        """Instantiate the WindowDataStorageFactory specified by this config."""
        from rslearn.dataset.window_data_storage.storage import (
            WindowDataStorageFactory,
        )
        from rslearn.utils.jsonargparse import init_jsonargparse

        init_jsonargparse()
        parser = jsonargparse.ArgumentParser()
        parser.add_argument("--wds", type=WindowDataStorageFactory)
        cfg = parser.parse_object(
            {
                "wds": dict(
                    class_path=self.class_path,
                    init_args=self.init_args,
                )
            }
        )
        return parser.instantiate_classes(cfg).wds

instantiate_factory

instantiate_factory() -> WindowDataStorageFactory

Instantiate the WindowDataStorageFactory specified by this config.

Source code in rslearn/config/dataset.py

def instantiate_factory(self) -> "WindowDataStorageFactory":
    """Instantiate the WindowDataStorageFactory specified by this config."""
    from rslearn.dataset.window_data_storage.storage import (
        WindowDataStorageFactory,
    )
    from rslearn.utils.jsonargparse import init_jsonargparse

    init_jsonargparse()
    parser = jsonargparse.ArgumentParser()
    parser.add_argument("--wds", type=WindowDataStorageFactory)
    cfg = parser.parse_object(
        {
            "wds": dict(
                class_path=self.class_path,
                init_args=self.init_args,
            )
        }
    )
    return parser.instantiate_classes(cfg).wds

DatasetConfig

Bases: BaseModel

Overall dataset configuration.

Source code in rslearn/config/dataset.py

class DatasetConfig(BaseModel):
    """Overall dataset configuration."""

    model_config = ConfigDict(extra="forbid")

    layers: dict[str, LayerConfig] = Field(description="Layers in the dataset.")
    tile_store: dict[str, Any] = Field(
        default={"class_path": "rslearn.tile_stores.default.DefaultTileStore"},
        description="jsonargparse configuration for the TileStore.",
    )
    storage: StorageConfig = Field(
        default_factory=lambda: StorageConfig(),
        description="jsonargparse configuration for the WindowStorageFactory.",
    )
    window_data_storage: WindowDataStorageConfig = Field(
        default_factory=lambda: WindowDataStorageConfig(),
        description=(
            "jsonargparse configuration for the WindowDataStorage. "
            "Controls how materialized raster/vector data is laid out on "
            "disk inside each window. Defaults to per-item-group layout, "
            "compatible with all existing rslearn datasets."
        ),
    )

    @field_validator("layers", mode="after")
    @classmethod
    def layer_names_validator(cls, v: dict[str, LayerConfig]) -> dict[str, LayerConfig]:
        """Ensure layer names don't contain periods, since we use periods to distinguish different materialized groups within a layer."""
        for layer_name in v.keys():
            if "." in layer_name:
                raise ValueError(f"layer names must not contain periods: {layer_name}")
        return v

layer_names_validator classmethod

layer_names_validator(v: dict[str, LayerConfig]) -> dict[str, LayerConfig]

Ensure layer names don't contain periods, since we use periods to distinguish different materialized groups within a layer.

Source code in rslearn/config/dataset.py

@field_validator("layers", mode="after")
@classmethod
def layer_names_validator(cls, v: dict[str, LayerConfig]) -> dict[str, LayerConfig]:
    """Ensure layer names don't contain periods, since we use periods to distinguish different materialized groups within a layer."""
    for layer_name in v.keys():
        if "." in layer_name:
            raise ValueError(f"layer names must not contain periods: {layer_name}")
    return v

ensure_timedelta

ensure_timedelta(v: Any) -> Any

Ensure the value is a timedelta.

If the value is a string, we try to parse it with pytimeparse.

This function is meant to be used like Annotated[timedelta, BeforeValidator(ensure_timedelta)].

Source code in rslearn/config/dataset.py

def ensure_timedelta(v: Any) -> Any:
    """Ensure the value is a timedelta.

    If the value is a string, we try to parse it with pytimeparse.

    This function is meant to be used like Annotated[timedelta, BeforeValidator(ensure_timedelta)].
    """
    if isinstance(v, timedelta):
        return v
    if isinstance(v, str):
        return pytimeparse.parse(v)
    raise TypeError(f"Invalid type for timedelta: {type(v).__name__}")

ensure_optional_timedelta

ensure_optional_timedelta(v: Any) -> Any

Like ensure_timedelta, but allows None as a value.

Source code in rslearn/config/dataset.py

def ensure_optional_timedelta(v: Any) -> Any:
    """Like ensure_timedelta, but allows None as a value."""
    if v is None:
        return None
    if isinstance(v, timedelta):
        return v
    if isinstance(v, str):
        return pytimeparse.parse(v)
    raise TypeError(f"Invalid type for timedelta: {type(v).__name__}")

serialize_optional_timedelta

serialize_optional_timedelta(v: timedelta | None) -> str | None

Serialize an optional timedelta for compatibility with pytimeparse.

Source code in rslearn/config/dataset.py

def serialize_optional_timedelta(v: timedelta | None) -> str | None:
    """Serialize an optional timedelta for compatibility with pytimeparse."""
    if v is None:
        return None
    return str(v.total_seconds()) + "s"