Source code for forecasting.simulation.harness_base

from typing import List, Dict, Optional, Union
import polars as pl

from forecasting.data.storage import Storage




class BaseHarness:
    """
    Abstract base class for all Harnesses.

    Provides common functionality for:
    - S3/Local/Modal volume file path resolution.
    - Cloud storage configuration.
    """

    def __init__(
        self,
        base_source: Union[str, List[str]],
        storage_options: Optional[Dict] = None,
        id: str = "harness",
    ):
        self.base_source = base_source
        self.storage_options = storage_options or {}
        self.id = id

    def _resolve_paths(self) -> List[str]:
        """Resolves wildcards and S3/Modal volume paths to a concrete list of files."""
        paths = []

        # Normalize to list
        sources = (
            self.base_source
            if isinstance(self.base_source, list)
            else [self.base_source]
        )

        for source in sources:
            if not isinstance(source, str):
                continue
            paths.extend(Storage.list_files(source))

        # Deduplicate and sort
        return sorted(list(set(paths)))