dsgrid-toolkit 0.3.3__cp313-cp313-win_amd64.whl

dsgrid/tests/common.py

@@ -0,0 +1,140 @@
+import getpass
+import os
+from pathlib import Path
+
+import pytest
+from semver import VersionInfo
+
+from dsgrid.exceptions import DSGInvalidParameter, DSGInvalidOperation
+from dsgrid.registry.dimension_registry_manager import DimensionRegistryManager
+from dsgrid.registry.project_registry_manager import ProjectRegistryManager
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.registry.common import DataStoreType, DatabaseConnection, VersionUpdateType
+from dsgrid.utils.files import dump_data, load_data
+
+TEST_PROJECT_PATH = Path(__file__).absolute().parents[2] / "dsgrid-test-data"
+TEST_PROJECT_REPO = TEST_PROJECT_PATH / "test_efs"
+TEST_STANDARD_SCENARIOS_PROJECT_REPO = TEST_PROJECT_PATH / "standard_scenarios_2021"
+TEST_DATASET_DIRECTORY = TEST_PROJECT_PATH / "datasets"
+TEST_REGISTRY_DATABASE = "cached-test-dsgrid"
+TEST_REGISTRY_BASE_PATH = Path("tests/data/registry")
+TEST_REGISTRY_DATA_PATH = Path("tests/data/registry/registry_data")
+TEST_EFS_REGISTRATION_FILE = Path("tests/data/test_efs_registration.json5")
+# AWS_PROFILE_NAME = "nrel-aws-dsgrid"
+TEST_REMOTE_REGISTRY = "s3://nrel-dsgrid-registry-test"
+CACHED_TEST_REGISTRY_DB = f"sqlite:///{TEST_REGISTRY_BASE_PATH}/cached_registry.db"
+STANDARD_SCENARIOS_PROJECT_REPO = Path(__file__).parents[2] / "dsgrid-project-StandardScenarios"
+IEF_PROJECT_REPO = Path(__file__).parents[2] / "dsgrid-project-IEF"
+SIMPLE_STANDARD_SCENARIOS = TEST_PROJECT_PATH / "filtered_registries" / "simple_standard_scenarios"
+SIMPLE_STANDARD_SCENARIOS_REGISTRY_DB = (
+    f"sqlite:///{TEST_PROJECT_PATH}/filtered_registries/simple_standard_scenarios/registry.db"
+)
+
+
+def create_local_test_registry(
+    tmpdir: Path, conn=None, data_store_type: DataStoreType = DataStoreType.FILESYSTEM
+):
+    if conn is None:
+        conn = DatabaseConnection(url=f"sqlite:///{tmpdir}/dsgrid-test.db")
+    data_path = tmpdir / "registry_data"
+    mgr = RegistryManager.create(conn, data_path, data_store_type=data_store_type, overwrite=True)
+    mgr.dispose()
+    return data_path
+
+
+def check_configs_update(base_dir, manager):
+    """Runs an update on one of each type of config.
+
+    Parameters
+    ----------
+    base_dir : Path
+    manager : RegistryManager
+
+    Returns
+    -------
+    list
+        Each updated config and new version: [(Updated ID, new version)]
+        For the dimension element the tuple is (Updated ID, dimension type, new version).
+        Order is dimension ID, dimension mapping ID, dataset ID, project ID
+
+    """
+    update_dir = base_dir / "updates"
+    user = getpass.getuser()
+
+    updated_ids = []
+    for mgr in (
+        manager.dimension_manager,
+        manager.dimension_mapping_manager,
+        manager.dataset_manager,
+        manager.project_manager,
+    ):
+        config_id = mgr.list_ids()[0]
+        version = mgr.get_latest_version(config_id)
+        check_config_update(update_dir, mgr, config_id, user, version)
+        new_version = mgr.get_latest_version(config_id)
+        if isinstance(mgr, DimensionRegistryManager):
+            config = mgr.get_by_id(config_id)
+            updated_ids.append((config_id, config.model.dimension_type, new_version))
+        else:
+            updated_ids.append((config_id, new_version))
+
+    return updated_ids
+
+
+def check_config_update(base_dir, mgr: ProjectRegistryManager, config_id, user, version):
+    """Runs basic positive and negative update tests for the config.
+
+    Parameters
+    ----------
+    base_dir : str
+    mgr : RegistryManagerBase
+    config_id : str
+    user : str
+    version : str
+
+    """
+    config_file = Path(base_dir) / mgr.config_class().config_filename()
+    assert not config_file.exists()
+    try:
+        mgr.dump(config_id, base_dir, force=True)
+        with pytest.raises(DSGInvalidOperation):
+            mgr.dump(config_id, base_dir)
+        mgr.dump(config_id, base_dir, force=True)
+        assert config_file.exists()
+        config_data = load_data(config_file)
+        config_data["description"] += "; updated description"
+        dump_data(config_data, config_file)
+        with pytest.raises(DSGInvalidParameter):
+            mgr.update_from_file(
+                config_file,
+                "invalid_config_id",
+                user,
+                VersionUpdateType.PATCH,
+                "update to description",
+                version,
+            )
+        with pytest.raises(DSGInvalidParameter):
+            mgr.update_from_file(
+                config_file,
+                config_id,
+                user,
+                VersionUpdateType.PATCH,
+                "update to description",
+                str(VersionInfo.parse(version).bump_patch()),
+            )
+
+        mgr.update_from_file(
+            config_file,
+            config_id,
+            user,
+            VersionUpdateType.PATCH,
+            "update to description",
+            version,
+        )
+        assert (
+            VersionInfo.parse(mgr.get_latest_version(config_id))
+            == VersionInfo.parse(version).bump_patch()
+        )
+    finally:
+        if config_file.exists():
+            os.remove(config_file)

dsgrid/tests/make_us_data_registry.py

@@ -0,0 +1,265 @@
+import contextlib
+import getpass
+import logging
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import rich_click as click
+
+from dsgrid.cli.common import path_callback
+from dsgrid.loggers import setup_logging, check_log_file_size
+from dsgrid.registry.common import DataStoreType, DatabaseConnection
+from dsgrid.registry.registry_manager import RegistryManager
+from dsgrid.tests.common import (
+    create_local_test_registry,
+    TEST_REMOTE_REGISTRY,
+    TEST_PROJECT_REPO,
+    TEST_DATASET_DIRECTORY,
+)
+from dsgrid.utils.timing import timer_stats_collector
+from dsgrid.utils.files import dump_data, load_data
+from dsgrid.utils.id_remappings import (
+    map_dimension_names_to_ids,
+    replace_dimension_names_with_current_ids,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def _find_file_with_stem(directory: Path, stem: str) -> Path | None:
+    """Find a file in directory with the given stem, regardless of extension."""
+    for path in directory.iterdir():
+        if path.stem == stem:
+            return path
+    return None
+
+
+def update_dataset_config_paths(config_file: Path, dataset_id: str) -> None:
+    """Update the data file paths in a dataset config to be relative to the config file.
+
+    Parameters
+    ----------
+    config_file : Path
+        Path to the dataset configuration file.
+    dataset_id : str
+        The dataset ID, used to locate the data files in TEST_DATASET_DIRECTORY.
+    """
+    data = load_data(config_file)
+    if "data_layout" not in data:
+        return
+
+    data_layout = data["data_layout"]
+    config_dir = config_file.parent.resolve()
+    dataset_data_dir = (TEST_DATASET_DIRECTORY / dataset_id).resolve()
+
+    if "data_file" in data_layout:
+        stem = Path(data_layout["data_file"]["path"]).stem
+        data_file_path = _find_file_with_stem(dataset_data_dir, stem)
+        if data_file_path is None:
+            msg = f"Could not find data file with stem '{stem}' in {dataset_data_dir}"
+            raise FileNotFoundError(msg)
+        relative_path = os.path.relpath(data_file_path, config_dir)
+        data_layout["data_file"]["path"] = relative_path
+
+    if "lookup_data_file" in data_layout and data_layout["lookup_data_file"] is not None:
+        stem = Path(data_layout["lookup_data_file"]["path"]).stem
+        lookup_file_path = _find_file_with_stem(dataset_data_dir, stem)
+        if lookup_file_path is None:
+            msg = f"Could not find lookup file with stem '{stem}' in {dataset_data_dir}"
+            raise FileNotFoundError(msg)
+        relative_path = os.path.relpath(lookup_file_path, config_dir)
+        data_layout["lookup_data_file"]["path"] = relative_path
+
+    if "missing_associations" in data_layout and data_layout["missing_associations"] is not None:
+        items = []
+        for item in data_layout["missing_associations"]:
+            stem = Path(item).stem
+            missing_path = _find_file_with_stem(dataset_data_dir, stem)
+            if missing_path is None:
+                msg = (
+                    f"Could not find missing associations with stem '{stem}' in {dataset_data_dir}"
+                )
+                raise FileNotFoundError(msg)
+            relative_path = os.path.relpath(missing_path, config_dir)
+            items.append(relative_path)
+        data_layout["missing_associations"] = items
+
+    dump_data(data, config_file)
+
+
+@contextlib.contextmanager
+def make_test_data_registry(
+    registry_path,
+    src_dir,
+    include_projects=True,
+    include_datasets=True,
+    offline_mode=True,
+    database_url: str | None = None,
+    data_store_type: DataStoreType = DataStoreType.FILESYSTEM,
+):
+    """Creates a local registry from a dsgrid project source directory for testing.
+
+    This is a context manager that yields the RegistryManager and disposes it on exit.
+
+    Parameters
+    ----------
+    registry_path : Path
+        Path in which the registry will be created.
+    src_dir : Path
+        Path containing source config files
+    include_projects : bool
+        If False, do not register any projects.
+    include_datasets : bool
+        If False, do not register any datasets.
+    offline_mode : bool
+        If False, use the test remote registry.
+    data_store_type: DataStoreType
+        Type of store to use for the registry data.
+
+    Yields
+    ------
+    RegistryManager
+    """
+    if not include_projects and include_datasets:
+        msg = "If include_datasets is True then include_projects must also be True."
+        raise Exception(msg)
+    url = f"sqlite:///{registry_path}/registry.db" if database_url is None else database_url
+    conn = DatabaseConnection(url=url)
+    create_local_test_registry(registry_path, conn=conn, data_store_type=data_store_type)
+    dataset_dirs = [
+        Path("datasets/modeled/comstock"),
+        Path("datasets/modeled/comstock_unpivoted"),
+    ]
+
+    user = getpass.getuser()
+    log_message = "Initial registration"
+    if offline_mode:
+        manager = RegistryManager.load(conn, offline_mode=offline_mode)
+    else:
+        manager = RegistryManager.load(
+            conn, remote_path=TEST_REMOTE_REGISTRY, offline_mode=offline_mode
+        )
+
+    try:
+        project_config_file = src_dir / "project.json5"
+        project_id = load_data(project_config_file)["project_id"]
+        dataset_config_files = [src_dir / path / "dataset.json5" for path in dataset_dirs]
+        dataset_mapping_files = [
+            src_dir / path / "dimension_mappings.json5" for path in dataset_dirs
+        ]
+        for i, filename in enumerate(dataset_mapping_files):
+            if not filename.exists():
+                dataset_mapping_files[i] = None
+        dataset_ids = [
+            load_data(config_file)["dataset_id"] for config_file in dataset_config_files
+        ]
+
+        if include_projects:
+            print("\n 1. register project: \n")
+            manager.project_manager.register(
+                project_config_file,
+                user,
+                log_message,
+            )
+        if include_datasets:
+            for i, dataset_config_file in enumerate(dataset_config_files):
+                dataset_id = dataset_ids[i]
+                print(f"\n 2. register dataset {dataset_id}: \n")
+                dataset_mapping_file = dataset_mapping_files[i]
+                mappings = map_dimension_names_to_ids(manager.dimension_manager)
+                replace_dimension_names_with_current_ids(dataset_config_file, mappings)
+                update_dataset_config_paths(dataset_config_file, dataset_id)
+                manager.dataset_manager.register(
+                    dataset_config_file,
+                    user,
+                    log_message,
+                )
+                print(f"\n 3. submit dataset {dataset_id} to project\n")
+                manager.project_manager.submit_dataset(
+                    project_id,
+                    dataset_id,
+                    user,
+                    log_message,
+                    dimension_mapping_file=dataset_mapping_file,
+                )
+        yield manager
+    finally:
+        manager.dispose()
+
+
+@click.command()
+@click.argument(
+    "registry-path",
+    type=Path,
+    default=f"{Path.home()}/.dsgrid-test-registry",
+    callback=path_callback,
+)
+@click.option(
+    "-f",
+    "--force",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="Delete registry-path if it exists.",
+)
+@click.option(
+    "-p",
+    "--project-dir",
+    default=TEST_PROJECT_REPO,
+    help="path to a project repository",
+    callback=path_callback,
+)
+@click.option(
+    "-t",
+    "--data-store-type",
+    type=click.Choice([x.value for x in DataStoreType]),
+    default=DataStoreType.FILESYSTEM.value,
+    show_default=True,
+    help="Type of store to use for the registry data.",
+    callback=lambda *x: DataStoreType(x[2]),
+)
+@click.option(
+    "--verbose",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help="Enable verbose log output.",
+)
+def run(
+    registry_path: Path,
+    force: bool,
+    project_dir: Path,
+    data_store_type: DataStoreType,
+    verbose: bool,
+):
+    """Creates a local registry from a dsgrid project source directory for testing."""
+    level = logging.DEBUG if verbose else logging.INFO
+    log_file = Path("test_dsgrid_project.log")
+    check_log_file_size(log_file, no_prompts=True)
+    setup_logging("dsgrid", log_file, console_level=level, file_level=level, mode="a")
+    if registry_path.exists():
+        if force:
+            shutil.rmtree(registry_path)
+        else:
+            print(f"{registry_path} already exists. Use --force to overwrite.")
+    os.makedirs(registry_path)
+    tmp_project_dir = Path(tempfile.gettempdir()) / "tmp_test_project_dir"
+    if tmp_project_dir.exists():
+        shutil.rmtree(tmp_project_dir)
+    shutil.copytree(project_dir, tmp_project_dir)
+    try:
+        with make_test_data_registry(
+            registry_path,
+            tmp_project_dir / "dsgrid_project",
+            data_store_type=data_store_type,
+        ):
+            pass  # Manager is created and disposed in context manager
+    finally:
+        timer_stats_collector.log_stats()
+
+
+if __name__ == "__main__":
+    run()

dsgrid/tests/register_derived_datasets.py

@@ -0,0 +1,103 @@
+import logging
+import shutil
+from pathlib import Path
+
+import rich_click as click
+
+from dsgrid.loggers import setup_logging, check_log_file_size
+from dsgrid.query.models import ProjectQueryModel
+from dsgrid.registry.dataset_registry import DatasetRegistry
+from dsgrid.utils.run_command import check_run_command
+from dsgrid.utils.timing import timer_stats_collector
+
+
+logger = logging.getLogger(__name__)
+
+
+@click.command()
+@click.argument(
+    "query_files",
+    nargs=-1,
+    type=click.Path(exists=True),
+    callback=lambda *x: [Path(y) for y in x[2]],
+)
+@click.option(
+    "-r", "--registry-path", required=True, callback=lambda *x: Path(x[2]), help="Path to registry"
+)
+@click.option(
+    "-o",
+    "--output",
+    default="query_output",
+    show_default=True,
+    type=click.Path(),
+    help="Output directory for query results",
+    callback=lambda *x: Path(x[2]),
+)
+@click.option(
+    "-p",
+    "--project-id",
+    default="dsgrid_conus_2022",
+    show_default=True,
+    type=str,
+    help="Project ID",
+)
+@click.option(
+    "--verbose", is_flag=True, default=False, show_default=True, help="Enable verbose log output."
+)
+def run(query_files, project_id, registry_path, output, verbose):
+    """Registers derived datasets in a local registry for testing."""
+    level = logging.DEBUG if verbose else logging.INFO
+    log_file = Path("dsgrid_registration.log")
+    check_log_file_size(log_file, no_prompts=True)
+    logger = setup_logging(__name__, log_file, console_level=level, file_level=level, mode="a")
+    try:
+        _run_registration(query_files, project_id, registry_path, output)
+    finally:
+        # Raise the console level so that timer stats only go to the log file.
+        for _, handler in enumerate(logger.handlers):
+            if handler.name == "console":
+                handler.setLevel(logging.WARNING)
+                break
+
+        timer_stats_collector.log_stats()
+
+
+def _run_registration(
+    query_files: list[Path], project_id: str, registry_path: Path, query_output_dir: Path
+):
+    log_message = "Submit derived dataset"
+    query_output_dir.mkdir(exist_ok=True)
+    derived_dataset_config_dir = query_output_dir / "derived_dataset_configs"
+    if derived_dataset_config_dir.exists():
+        shutil.rmtree(derived_dataset_config_dir)
+    derived_dataset_config_dir.mkdir()
+    for query_file in query_files:
+        logger.info("Register derived dataset from %s", query_file)
+        query = ProjectQueryModel.from_file(query_file)
+        dataset_id = query.project.dataset.dataset_id
+        dataset_config_dir = derived_dataset_config_dir / dataset_id
+        dataset_config_file = dataset_config_dir / DatasetRegistry.config_filename()
+
+        create_cmd = (
+            f"dsgrid query project run --registry-path={registry_path} "
+            f"-o {query_output_dir} {query_file}"
+        )
+
+        config_cmd = (
+            f"dsgrid query project create-derived-dataset-config "
+            f"--registry-path={registry_path} {query_output_dir / dataset_id} {dataset_config_dir}"
+        )
+
+        submit_cmd = (
+            f"dsgrid registry --path {registry_path} projects "
+            f"register-and-submit-dataset -c {dataset_config_file} -p {project_id} "
+            f"-l '{log_message}' -d {query_output_dir / dataset_id}"
+        )
+
+        for cmd in (create_cmd, config_cmd, submit_cmd):
+            logger.info(cmd)
+            check_run_command(cmd)
+
+
+if __name__ == "__main__":
+    run()

dsgrid/tests/utils.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+
+from dsgrid.spark.types import use_duckdb
+from dsgrid.utils.spark import get_spark_session
+
+
+def read_parquet(filename: Path):
+    """Read a Parquet file and load it into cache. This helps debugging with pytest --pdb.
+    If you don't use this, the parquet file will get deleted on a failure and you won't be able
+    to inspect the dataframe.
+    """
+    spark = get_spark_session()
+    df = spark.read.parquet(str(filename))
+    if not use_duckdb():
+        df.cache()
+        df.count()
+    return df
+
+
+def read_parquet_two_table_format(path: Path):
+    spark = get_spark_session()
+    load_data = spark.read.parquet(str(path / "load_data.parquet"))
+    lookup = spark.read.parquet(str(path / "load_data_lookup.parquet"))
+    table = load_data.join(lookup, on="id").drop("id")
+    return table

dsgrid/time/time_conversions.py

@@ -0,0 +1,80 @@
+"""Functions to perform time conversions"""
+
+from datetime import datetime
+
+from dsgrid.time.types import DayType, Season
+
+
+def convert_datetime_to_day_type(timestamp):
+    """Returns the day type for the datetime.
+
+    Parameters
+    ----------
+    timestamp : datetime.datetime
+
+    Returns
+    -------
+    str
+        DayType id
+
+    """
+    # Monday is 0.
+    if timestamp.weekday() <= 4:
+        return DayType.WEEKDAY.value
+    return DayType.WEEKEND.value
+
+
+def convert_datetime_to_season(timestamp):
+    """Returns the season for the datetime.
+
+    Parameters
+    ----------
+    timestamp : datetime.datetime
+
+    Returns
+    -------
+    str
+        Season id
+
+    """
+    # TODO: dates do change slightly every year. Is this close enough?
+    # dates also change by region, it's weather driven.
+    year = timestamp.year
+    if timestamp < datetime(year, 3, 20) or timestamp > datetime(year, 12, 21):
+        season = Season.WINTER.value
+    elif timestamp < datetime(year, 6, 20):
+        season = Season.SPRING.value
+    elif timestamp < datetime(year, 9, 22):
+        season = Season.SUMMER.value
+    else:
+        season = Season.AUTUMN.value
+
+    return season
+
+
+def interpret_datetime(timestamp):
+    """Return a datetime object from a timestamp string.
+
+    Parameters
+    ----------
+    timestamp : str
+
+    Returns
+    -------
+    datetime.datetime
+
+    """
+    formats = (
+        "%Y-%m-%dT%H:%M:%S",
+        "%Y-%m-%d %H:%M:%S",
+        "%Y-%m-%dT%H:%M:%SZ",
+        "%Y-%m-%dT%H:%M:%S.%f",
+        "%Y-%m-%d %H:%M:%S.%f",
+    )
+
+    for i, fmt in enumerate(formats):
+        try:
+            return datetime.strptime(timestamp, fmt)
+        except ValueError:
+            if i == len(formats) - 1:
+                raise

dsgrid/time/types.py

@@ -0,0 +1,67 @@
+"""Types related to time"""
+
+from datetime import datetime
+from typing import NamedTuple
+
+from dsgrid.data_models import DSGEnum
+
+
+class DayType(DSGEnum):
+    """Day types"""
+
+    WEEKEND = "weekend"
+    WEEKDAY = "weekday"
+
+
+class Season(DSGEnum):
+    """Seasons"""
+
+    WINTER = "winter"
+    SPRING = "spring"
+    SUMMER = "summer"
+    AUTUMN = "autumn"
+    FALL = "autumn"
+
+
+# The types below represent the timestamps that exist as columns in all datasets.
+
+
+class DatetimeTimestampType(NamedTuple):
+    """Single column with datetime."""
+
+    timestamp: datetime
+
+
+class AnnualTimestampType(NamedTuple):
+    """Single column with only year."""
+
+    time_year: int
+
+
+class OneWeekPerMonthByHourType(NamedTuple):
+    """Columns of representative time with one week per month."""
+
+    month: int
+    # 0 = Monday, 6 = Sunday. Follows pyspark.sql.functions.weekday and Python datetime.weekday.
+    day_of_week: int
+    hour: int
+
+
+class OneWeekdayDayAndOneWeekendDayPerMonthByHourType(NamedTuple):
+    """Columns of representative time with month, hour, and weekday vs weekend."""
+
+    month: int
+    is_weekday: bool
+    hour: int
+
+
+class IndexTimestampType(NamedTuple):
+    """Single column with numerical indices."""
+
+    time_index: int
+
+
+class StringTimestampType(NamedTuple):
+    """Single column with time (must include offset) as str."""
+
+    timestamp: str