Source code for rastervision.pipeline.pipeline_config

from typing import TYPE_CHECKING
from os.path import join

from rastervision.pipeline.config import Config, Field
from rastervision.pipeline.config import register_config

if TYPE_CHECKING:
    from typing import Self
    from rastervision.pipeline.pipeline import Pipeline


@register_config('pipeline')
class PipelineConfig(Config):
    """Base class for configuring :class:`Pipelines <.Pipeline>`.

    This should be subclassed to configure new Pipelines.
    """
    root_uri: str | None = Field(
        None, description='The root URI for output generated by the pipeline')
    rv_config: dict | None = Field(
        None,
        description='Used to store serialized RVConfig so pipeline can '
        'run in remote environment with the local RVConfig. This should '
        'not be set explicitly by users -- it is only used by the runner '
        'when running a remote pipeline.')
    plugin_versions: dict[str, int] | None = Field(
        None,
        description=
        ('Used to store a mapping of plugin module paths to the latest '
         'version number. This should not be set explicitly by users -- it is set '
         'automatically when serializing and saving the config to disk.'))

    def get_config_uri(self) -> str:
        """Get URI of serialized version of this PipelineConfig."""
        return join(self.root_uri, 'pipeline-config.json')

    def build(self, tmp_dir: str) -> 'Pipeline':
        """Return a pipeline based on this configuration.

        Subclasses should override this to return an instance of the
        corresponding subclass of Pipeline.

        Args:
            tmp_dir: root of any temporary directory to pass to pipeline
        """
        from rastervision.pipeline.pipeline import Pipeline  # noqa
        return Pipeline(self, tmp_dir)

    def dict(self, with_rv_metadata: bool = True, **kwargs) -> dict:
        # override to have with_rv_metadata=True by default
        cfg_dict = super().dict(with_rv_metadata=with_rv_metadata, **kwargs)
        return cfg_dict

    @classmethod
    def from_dict(cls, cfg_dict: dict) -> 'Self':
        # override to retain plugin_versions
        from rastervision.pipeline.config import build_config, upgrade_config
        cfg_dict: dict = upgrade_config(cfg_dict)
        cfg = build_config(cfg_dict)
        return cfg