test for _check_arg

relative

Author: Borda

setup.cfg

@@ -12,6 +12,7 @@ addopts =
 [coverage:run]
 parallel = True
 concurrency = thread
+relative_files = True
 
 [coverage:report]
 exclude_lines =

tests/bases/test_collections.py

@@ -182,3 +182,11 @@ def test_metric_collection_same_order():
     col2 = MetricCollection({"b": m2, "a": m1})
     for k1, k2 in zip(col1.keys(), col2.keys()):
         assert k1 == k2
+
+
+def test_collection_check_arg():
+    assert MetricCollection._check_arg(None, 'prefix') is None
+    assert MetricCollection._check_arg('sample', 'prefix') == 'sample'
+
+    with pytest.raises(ValueError, match="Expected input `postfix` to be a string, but got"):
+        MetricCollection._check_arg(1, 'postfix')

torchmetrics/collections.py

@@ -191,9 +191,7 @@ def _set_name(self, base: str) -> str:
         return name
 
     @staticmethod
-    def _check_arg(arg: str, name: str) -> Optional[str]:
-        if arg is not None:
-            if isinstance(arg, str):
-                return arg
-            raise ValueError(f'Expected input {name} to be a string')
-        return None
+    def _check_arg(arg: Optional[str], name: str) -> Optional[str]:
+        if arg is None or isinstance(arg, str):
+            return arg
+        raise ValueError(f'Expected input `{name}` to be a string, but got {type(arg)}')