Skip to content

Commit

Permalink
Add .sdmx.quantity_to_message(), tests
Browse files Browse the repository at this point in the history
  • Loading branch information
khaeru committed Jan 26, 2024
1 parent bc27e81 commit bc8ac26
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 8 deletions.
19 changes: 11 additions & 8 deletions genno/compat/sdmx.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Dict, Iterable, List, Mapping, Optional, Union
from typing import Dict, Iterable, List, Mapping, Optional, Union

from genno import Quantity

Expand All @@ -10,9 +10,6 @@
HAS_SDMX = True


if TYPE_CHECKING:
import sdmx.model.common

__all__ = [
"codelist_to_groups",
"dataset_to_quantity",
Expand Down Expand Up @@ -94,10 +91,7 @@ def quantity_to_dataset(
) -> "sdmx.model.v21.Dataset":
"""Convert :class:`.Quantity to :class:`DataSet <sdmx.model.common.BaseDataSet>`.
Currently:
- `structure` must be provided.
- The resulting data set is structure-specific and flat (not grouped into Series).
The resulting data set is structure-specific and flat (not grouped into Series).
"""
try:
# URN of DSD stored on `qty` matches `structure`
Expand All @@ -124,3 +118,12 @@ def as_obs(key, value):
ds.obs.extend(as_obs(key, value) for key, value in qty.to_series().items())

return ds


def quantity_to_message(
qty: Quantity, structure: "sdmx.model.v21.DataStructureDefinition", **kwargs
) -> "sdmx.message.DataMessage":
"""Convert :class:`.Quantity to :class:`DataMessage <sdmx.message.DataMessage>`."""
dm = sdmx.message.DataMessage(**kwargs)
dm.data.append(quantity_to_dataset(qty, structure))
return dm
15 changes: 15 additions & 0 deletions genno/tests/compat/test_sdmx.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
codelist_to_groups,
dataset_to_quantity,
quantity_to_dataset,
quantity_to_message,
)
from genno.testing import add_test_data

Expand Down Expand Up @@ -91,3 +92,17 @@ def test_quantity_to_dataset(dsd, dm):

# Dataset is associated with its DSD
assert dsd is ds.structured_by


def test_quantity_to_message(dsd, dm):
ds = dm.data[0]
qty = dataset_to_quantity(ds)

header = dm.header

result = quantity_to_message(qty, structure=dsd, header=header)

# Currently False because `result.observation_dimension` is not set
with pytest.raises(AssertionError):
# Resulting message compares equal to the original ("round trip")
assert dm.compare(result)

0 comments on commit bc8ac26

Please sign in to comment.