Forbid reserved keys in Image

fepegar · Apr 23, 2020 · b162d2f · b162d2f
1 parent d28137d
commit b162d2f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/tests/data/test_image.py b/tests/data/test_image.py
@@ -16,3 +16,7 @@ def test_image_not_found(self):
     def test_wrong_path_type(self):
         with self.assertRaises(TypeError):
             Image(5, INTENSITY)
+
+    def test_bad_key(self):
+        with self.assertRaises(ValueError):
+            Image(5, INTENSITY, affine=1)
diff --git a/torchio/data/image.py b/torchio/data/image.py
@@ -9,7 +9,7 @@
 import torch
 import numpy as np
 
-from ..torchio import TypePath
+from ..torchio import TypePath, DATA, TYPE, AFFINE, PATH, STEM
 from .io import read_image
 
 
@@ -29,6 +29,10 @@ class Image(dict):
     """
 
     def __init__(self, path: TypePath, type_: str, **kwargs: Dict[str, Any]):
+        for key in (DATA, AFFINE, TYPE, PATH, STEM):
+            if key in kwargs:
+                raise ValueError(f'Key {key} is reserved. Use a different one')
+
         super().__init__(**kwargs)
         self.path = self._parse_path(path)
         self.type = type_