Add test for Python 3.10+ unions

Author: nevillelyh

python/coglet/adt.py

@@ -13,9 +13,7 @@
 def _is_union(tpe: type) -> bool:
     if typing.get_origin(tpe) is Union:
         return True
-    if sys.version_info[0] > 3 or (
-        sys.version_info[0] == 3 and sys.version_info[1] >= 10
-    ):
+    if sys.version_info >= (3, 10):
         from types import UnionType
 
         if typing.get_origin(tpe) is UnionType:
@@ -144,7 +142,7 @@ def from_type(tpe: type):
             assert len(t_args) == 2 and type(None) in t_args, (
                 f'unsupported union type {tpe}'
             )
-            elem_t = t_args[0] if t_args[1] is type(None) else t_args[0]
+            elem_t = t_args[0] if t_args[1] is type(None) else t_args[1]
             # Fail fast to avoid the cryptic "unsupported Cog type" error later with elem_t
             nested_t = typing.get_origin(elem_t)
             assert nested_t is None, f'Optional cannot have nested type {nested_t}'

python/tests/runners/unions.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+from cog import BasePredictor
+from tests.util import check_python_version
+
+check_python_version(min_version=(3, 10))
+
+FIXTURE = [
+    (
+        {
+            'os1': 'foo0',
+            'os2': 'bar0',
+            'os3': 'baz0',
+        },
+        'foo0-bar0-baz0',
+    ),
+]
+
+
+class Predictor(BasePredictor):
+    test_inputs = {
+        'os1': 'foo',
+        'os2': 'bar',
+        'os3': 'baz',
+    }
+    setup_done = False
+
+    def setup(self) -> None:
+        self.setup_done = True
+
+    def predict(
+        self,
+        os1: Optional[str],
+        os2: str | None,
+        os3: None | str,
+    ) -> str:
+        return f'{os1}-{os2}-{os3}'

python/tests/test_test_inputs.py

@@ -5,6 +5,7 @@
 import pytest
 
 from coglet import inspector, runner, scope
+from tests.util import PythonVersionError
 
 
 def get_predictors() -> List[str]:
@@ -21,9 +22,12 @@ async def test_test_inputs(predictor):
     if predictor.startswith('function_'):
         entrypoint = 'predict'
 
-    p = inspector.create_predictor(module_name, entrypoint)
-    r = runner.Runner(p)
+    try:
+        p = inspector.create_predictor(module_name, entrypoint)
+        r = runner.Runner(p)
 
-    # Some predictors calls current_scope() and requires ctx_pid
-    scope.ctx_pid.set(predictor)
-    assert await r.test()
+        # Some predictors calls current_scope() and requires ctx_pid
+        scope.ctx_pid.set(predictor)
+        assert await r.test()
+    except PythonVersionError as e:
+        pytest.skip(reason=str(e))

python/tests/util.py

@@ -0,0 +1,20 @@
+import sys
+from typing import Optional, Tuple
+
+
+def check_python_version(
+    min_version: Optional[Tuple[int, int]] = None,
+    max_version: Optional[Tuple[int, int]] = None,
+) -> None:
+    if min_version is not None and sys.version_info < min_version:
+        raise PythonVersionError(
+            f'Python version must be >= {min_version[0]}.{min_version[1]}'
+        )
+    if max_version is not None and sys.version_info > max_version:
+        raise PythonVersionError(
+            f'Python version must be <= {max_version[0]}.{max_version[1]}'
+        )
+
+
+class PythonVersionError(Exception):
+    pass