BaseStorage.get() functionality (#5240)

* update

* changelog

Author: rusty1s

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.2.0] - 2022-MM-DD
 ### Added
+- Added `BaseStorage.get()` functionality ([#5240](https://github.com/pyg-team/pytorch_geometric/pull/5240))
 - Added a test to confirm that `to_hetero` works with `SparseTensor` ([#5222](https://github.com/pyg-team/pytorch_geometric/pull/5222))
 ### Changed
 ### Removed

test/data/test_data.py

@@ -22,6 +22,9 @@ def test_data():
 
     assert data.x.tolist() == x.tolist()
     assert data['x'].tolist() == x.tolist()
+    assert data.get('x').tolist() == x.tolist()
+    assert data.get('y', 2) == 2
+    assert data.get('y', None) is None
 
     assert sorted(data.keys) == ['edge_index', 'x']
     assert len(data) == 2

test/data/test_storage.py

@@ -13,6 +13,11 @@ def test_base_storage():
     assert storage.x is not None
     assert storage.y is not None
 
+    assert torch.allclose(storage.get('x', None), storage.x)
+    assert torch.allclose(storage.get('y', None), storage.y)
+    assert storage.get('z', 2) == 2
+    assert storage.get('z', None) is None
+
     assert len(list(storage.keys('x', 'y', 'z'))) == 2
     assert len(list(storage.keys('x', 'y', 'z'))) == 2
     assert len(list(storage.values('x', 'y', 'z'))) == 2

torch_geometric/data/storage.py

@@ -161,6 +161,9 @@ def apply(self, func: Callable, *args: List[str]):
 
     # Additional functionality ################################################
 
+    def get(self, key: str, value: Optional[Any] = None) -> Any:
+        return self._mapping.get(key, value)
+
     def to_dict(self) -> Dict[str, Any]:
         r"""Returns a dictionary of stored key/value pairs."""
         return copy.copy(self._mapping)