Function create_concat_datasets

xarray/tests/test_concat.py:45–83

(
    num_datasets: int = 2, seed: int | None = None, include_day: bool = True
)

# helper method to create multiple tests datasets to concat
def create_concat_datasets(
    num_datasets: int = 2, seed: int | None = None, include_day: bool = True
) -> list[Dataset]:
    rng = np.random.default_rng(seed)
    lat = rng.standard_normal(size=(1, 4))
    lon = rng.standard_normal(size=(1, 4))
    result = []
    variables = ["temperature", "pressure", "humidity", "precipitation", "cloud_cover"]
    for i in range(num_datasets):
        if include_day:
            data_tuple = (
                ["x", "y", "day"],
                rng.standard_normal(size=(1, 4, 2)),
            )
            data_vars = dict.fromkeys(variables, data_tuple)
            result.append(
                Dataset(
                    data_vars=data_vars,
                    coords={
                        "lat": (["x", "y"], lat),
                        "lon": (["x", "y"], lon),
                        "day": ["day" + str(i * 2 + 1), "day" + str(i * 2 + 2)],
                    },
                )
            )
        else:
            data_tuple = (
                ["x", "y"],
                rng.standard_normal(size=(1, 4)),
            )
            data_vars = dict.fromkeys(variables, data_tuple)
            result.append(
                Dataset(
                    data_vars=data_vars,
                    coords={"lat": (["x", "y"], lat), "lon": (["x", "y"], lon)},
                )
            )

    return result


# helper method to create multiple tests datasets to concat with specific types