Repository: cogent3/cogent3
Pull request #1202 (2022-11-16 05:58:54)

Restrict type hint complexity on composed apps


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index dca0e379c..144677b94 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -15,7 +15,7 @@ from typing import Tuple
 
 from scitrack import CachingLogger
 
-from cogent3.app.typing import get_constraint_names
+from cogent3.app.typing import get_constraint_names, type_tree
 from cogent3.util import parallel as PAR
 from cogent3.util import progress_display as UI
 from cogent3.util.misc import (
@@ -23,6 +23,7 @@ from cogent3.util.misc import (
     get_object_provenance,
     in_jupyter,
 )
+from cogent3.util.warning import discontinued
 
 from .data_store import (
     IGNORE,
@@ -44,7 +45,8 @@ __maintainer__ = "Gavin Huttley"
 __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Alpha"
 
-from ..util.warning import discontinued
+
+_builtin_seqs = list, set, tuple
 
 
 def _make_logfile_name(process):
@@ -802,6 +804,20 @@ def _get_raw_hints(main_func, min_params):
     if return_type is None:
         raise TypeError("NoneType invalid type for return value")
 
+    # we disallow type hints with too many levels of testing,
+    # e.g set[int] is ok but tuple[set[int]] is not
+    msg = (
+        "{} type {} nesting level exceeds 2 for {}"
+        "we suggest using a custom type, e.g. a dataclass"
+    )
+    depth, _ = type_tree(first_param_type)
+    if depth > 2:
+        raise TypeError(msg.format("first_param", first_param_type, depth))
+
+    depth, _ = type_tree(return_type)
+    if depth > 2:
+        raise TypeError(msg.format("return_type", return_type, depth))
+
     return first_param_type, return_type
 
 
@@ -962,9 +978,9 @@ def _validate_data_type(self, data):
     }:
         return True
 
-    if isinstance(data, (list, tuple)) and len(data):
-        data = data[0]
-    elif isinstance(data, (list, tuple)):
+    if isinstance(data, _builtin_seqs) and len(data):
+        data = next(iter(data))
+    elif isinstance(data, _builtin_seqs):
         return NotCompleted("ERROR", self, message=f"empty data", source=data)
 
     class_name = data.__class__.__name__
@@ -1162,7 +1178,7 @@ def define_app(klass=None, *, app_type: AppType = GENERIC):
             func.__name__ = meth
             setattr(klass, meth, func)
 
-        # Get and type hints of main function in klass
+        # Get type hints of main function in klass
         arg_hints, return_hint = _get_main_hints(klass)
         setattr(klass, "_data_types", arg_hints)
         setattr(klass, "_return_types", return_hint)
@@ -1260,7 +1276,7 @@ def _apply_to(
     if self.app_type is not LOADER and self.input:
         # As we will be explicitly calling the input object, we disconnect
         # the two-way interaction between input and self. This means self
-        # is not called twice, and self is not unecessarily pickled during
+        # is not called twice, and self is not unnecessarily pickled during
         # parallel execution.
         self.input = None
 
diff --git a/src/cogent3/app/typing.py b/src/cogent3/app/typing.py
index 11d175a69..0a8d607a3 100644
--- a/src/cogent3/app/typing.py
+++ b/src/cogent3/app/typing.py
@@ -121,3 +121,28 @@ def hints_from_strings(*strings: Iterable[str]) -> list:
             raise ValueError(f"{string!r} not a known type constant")
         types.append(_mappings[string])
     return types
+
+
+def type_tree(hint, depth=0) -> tuple:
+    """compute the order of types"""
+    level_type = get_origin(hint)
+    if not level_type:
+        return depth + 1, hint
+
+    levels = []
+    depths = []
+    for arg in get_args(hint):
+        d, t = type_tree(arg, depth=depth)
+        levels.append(t)
+        depths.append(d)
+    depth = max(depths) + 1
+
+    if len(levels) == 1:
+        levels = levels[0]
+
+    try:
+        levels = tuple(levels)
+    except TypeError:
+        levels = (levels,)
+
+    return depth, (level_type, levels)


Test patch:
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index 5658963fa..fd3bb4e9a 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -5,6 +5,7 @@ import pickle
 
 from pickle import dumps, loads
 from tempfile import TemporaryDirectory
+from typing import Set, Tuple
 from unittest import main
 from unittest.mock import Mock
 
@@ -1000,5 +1001,38 @@ def test_validate_data_type_not_completed_pass_through():
     __app_registry.pop(get_object_provenance(take_int2), None)
 
 
+@pytest.mark.parametrize("first,ret", ((Tuple[Set[str]], int), (int, Tuple[Set[str]])))
+def test_complex_type(first, ret):
+    # disallow >2-deep nesting of types for first arg and return type
+    with pytest.raises(TypeError):
+
+        @define_app
+        class x:
+            def main(self, data: first) -> ret:
+                return data
+
+
+@pytest.mark.parametrize("hint", (Tuple[Set[str]], Tuple[Tuple[Set[str]]]))
+def test_complex_type_depths(hint):
+    # disallow >2-deep nesting of types for first arg and return type
+    with pytest.raises(TypeError):
+
+        @define_app
+        class x:
+            def main(self, data: hint) -> bool:
+                return True
+
+
+@pytest.mark.parametrize("hint", (int, Set[str]))
+def test_complex_type_allowed_depths(hint):
+    # allowed <=2-deep nesting of types
+    @define_app
+    class x:
+        def main(self, data: hint) -> int:
+            return int
+
+    __app_registry.pop(get_object_provenance(x), None)
+
+
 if __name__ == "__main__":
     main()
diff --git a/tests/test_app/test_typing.py b/tests/test_app/test_typing.py
index ae41aa2be..ae68463b5 100644
--- a/tests/test_app/test_typing.py
+++ b/tests/test_app/test_typing.py
@@ -13,6 +13,7 @@ from cogent3.app.typing import (
     UnalignedSeqsType,
     get_constraint_names,
     hints_from_strings,
+    type_tree,
 )
 
 
@@ -119,3 +120,24 @@ def test_hint_inherited_class():
 
     got = get_constraint_names(dummy)
     assert got == frozenset(["dummy"])
+
+
+@pytest.mark.parametrize(
+    "hint,expect", ((int, 1), (Set[int], 2), (List[List[Set[float]]], 4))
+)
+def test_typing_tree_depth(hint, expect):
+    d, _ = type_tree(hint)
+    assert d == expect, (d, expect)
+
+
+@pytest.mark.parametrize(
+    "hint,expect",
+    (
+        (int, int),
+        (Set[int], (set, (int,))),
+        (List[Set[int]], (list, (set, (int,)))),
+    ),
+)
+def test_type_tree(hint, expect):
+    _, t = type_tree(hint)
+    assert t == expect, (t, expect)