Fix attribute keyword typo in save function (#35)

* Fix attribute keyword typo in save function and add test for it

* Add changelog entry

CHANGELOG.md

@@ -47,6 +47,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   [\#19](https://github.com/mllam/weather-model-graphs/pull/19)
   @joeloskarsson
 
+### Fixed
+
+- Fix `attribute` keyword bug in save function
+  [\#35](https://github.com/mllam/weather-model-graphs/pull/35)
+  @joeloskarsson
+
 ### Maintenance
 
 - Ensure that cell execution doesn't time out when building jupyterbook based

src/weather_model_graphs/save.py

@@ -108,7 +108,7 @@ def _concat_pyg_features(
         try:
             sub_graphs = list(
                 split_graph_by_edge_attribute(
-                    graph=graph, attribute=list_from_attribute
+                    graph=graph, attr=list_from_attribute
                 ).values()
             )
         except MissingEdgeAttributeError:

tests/test_save.py

@@ -1,6 +1,7 @@
 import tempfile
 
 import numpy as np
+import pytest
 from loguru import logger
 
 import weather_model_graphs as wmg
@@ -15,7 +16,8 @@ def _create_fake_xy(N=10):
     return xy
 
 
-def test_save_to_pyg():
+@pytest.mark.parametrize("list_from_attribute", [None, "level"])
+def test_save_to_pyg(list_from_attribute):
     if not HAS_PYG:
         logger.warning(
             "Skipping test_save_to_pyg because weather-model-graphs[pytorch] is not installed."
@@ -40,4 +42,9 @@ def test_save_to_pyg():
 
     with tempfile.TemporaryDirectory() as tmpdir:
         for name, graph in graph_components.items():
-            wmg.save.to_pyg(graph=graph, output_directory=tmpdir, name=name)
+            wmg.save.to_pyg(
+                graph=graph,
+                output_directory=tmpdir,
+                name=name,
+                list_from_attribute=list_from_attribute,
+            )