Add argsparse function for string input

brainglobe_utils/general/numerical.py

@@ -1,4 +1,5 @@
 import argparse
+from typing import Literal
 
 
 def is_even(num):
@@ -31,75 +32,82 @@ def is_even(num):
         return True
 
 
-def check_positive_float(value, none_allowed=True):
+def check_positive_float(
+    value: float | None | Literal["None", "none"],
+    none_allowed: bool = True,
+) -> float | None:
     """
     Used in argparse to enforce positive floats.
     Source: https://stackoverflow.com/questions/14117415
 
     Parameters
     ----------
-    value : float
+    value : float or None
         Input value.
 
     none_allowed : bool, optional
         If False, throw an error for None values.
 
     Returns
     -------
-    float
-        Input value, if it's positive.
+    float or None
+        Input value, if it's positive, or None.
 
     Raises
     ------
     argparse.ArgumentTypeError
         If input value is invalid.
     """
     ivalue = value
-    if ivalue is not None:
-        ivalue = float(ivalue)
-        if ivalue < 0:
+    if value in (None, "None", "none"):
+        if not none_allowed:
+            raise argparse.ArgumentTypeError(f"{ivalue} is an invalid value.")
+        value = None
+    else:
+        value = float(value)
+        if value < 0:
             raise argparse.ArgumentTypeError(
-                "%s is an invalid positive value" % value
+                f"{ivalue} is an invalid positive value"
             )
-    else:
-        if not none_allowed:
-            raise argparse.ArgumentTypeError("%s is an invalid value." % value)
 
-    return ivalue
+    return value
 
 
-def check_positive_int(value, none_allowed=True):
+def check_positive_int(
+    value: int | None | Literal["None", "none"], none_allowed: bool = True
+) -> int | None:
     """
     Used in argparse to enforce positive ints.
     Source: https://stackoverflow.com/questions/14117415
 
     Parameters
     ----------
-    value : int
+    value : int or None
         Input value.
 
     none_allowed : bool, optional
         If False, throw an error for None values.
 
     Returns
     -------
-    int
-        Input value, if it's positive.
+    int or None
+        Input value, if it's positive, or None.
 
     Raises
     ------
     argparse.ArgumentTypeError
         If input value is invalid.
     """
     ivalue = value
-    if ivalue is not None:
-        ivalue = int(ivalue)
-        if ivalue < 0:
+    if value in (None, "None", "none"):
+        if not none_allowed:
+            raise argparse.ArgumentTypeError(f"{ivalue} is an invalid value.")
+        value = None
+    else:
+        value = int(value)
+        if value < 0:
             raise argparse.ArgumentTypeError(
-                "%s is an invalid positive value" % value
+                f"{ivalue} is an invalid positive value"
             )
-    else:
-        if not none_allowed:
-            raise argparse.ArgumentTypeError("%s is an invalid value." % value)
 
-    return ivalue
+    return value

brainglobe_utils/general/string.py

@@ -1,5 +1,6 @@
+import argparse
 from pathlib import Path
-from typing import Optional
+from typing import Literal, Optional
 
 from natsort import natsorted
 
@@ -53,3 +54,38 @@ def get_text_lines(
     if return_lines is not None:
         lines = lines[return_lines]
     return lines
+
+
+def check_str(
+    value: str | None | Literal["None", "none"], none_allowed: bool = True
+) -> str | None:
+    """
+    Used in argparse to enforce str input.
+
+    Parameters
+    ----------
+    value : str or None
+        Input value.
+
+    none_allowed : bool, optional
+        If False, throw an error for None values.
+
+    Returns
+    -------
+    str or None
+        Input value, if it's str, or None.
+
+    Raises
+    ------
+    argparse.ArgumentTypeError
+        If input value is invalid.
+    """
+    ivalue = value
+    if value in (None, "None", "none"):
+        if not none_allowed:
+            raise argparse.ArgumentTypeError(f"{ivalue} is an invalid value.")
+        value = None
+    else:
+        value = str(value)
+
+    return value

tests/tests/test_general/test_numerical.py

@@ -25,10 +25,10 @@ def test_check_positive_float():
     with pytest.raises(ArgumentTypeError):
         assert numerical.check_positive_float(neg_val)
 
-    assert numerical.check_positive_float(None) is None
-
-    with pytest.raises(ArgumentTypeError):
-        assert numerical.check_positive_float(None, none_allowed=False)
+    for none_val in [None, "None", "none"]:
+        assert numerical.check_positive_float(none_val) is None
+        with pytest.raises(ArgumentTypeError):
+            assert numerical.check_positive_float(none_val, none_allowed=False)
 
     assert numerical.check_positive_float(0) == 0
 
@@ -42,9 +42,9 @@ def test_check_positive_int():
     with pytest.raises(ArgumentTypeError):
         assert numerical.check_positive_int(neg_val)
 
-    assert numerical.check_positive_int(None) is None
-
-    with pytest.raises(ArgumentTypeError):
-        assert numerical.check_positive_int(None, none_allowed=False)
+    for none_val in [None, "None", "none"]:
+        assert numerical.check_positive_int(none_val) is None
+        with pytest.raises(ArgumentTypeError):
+            assert numerical.check_positive_int(none_val, none_allowed=False)
 
     assert numerical.check_positive_int(0) == 0

tests/tests/test_general/test_string.py

@@ -1,3 +1,5 @@
+from argparse import ArgumentTypeError
+
 import pytest
 
 from brainglobe_utils.general import string
@@ -48,3 +50,13 @@ def test_get_string_lines(jabberwocky, jabberwocky_list):
         string.get_text_lines(jabberwocky, return_lines=8)
         == jabberwocky_list[8]
     )
+
+
+def test_check_str():
+    assert "me" == string.check_str("me")
+    assert "12" == string.check_str("12")
+
+    for none_val in [None, "None", "none"]:
+        assert string.check_str(none_val) is None
+        with pytest.raises(ArgumentTypeError):
+            assert string.check_str(none_val, none_allowed=False)