feature[next]: add support for Python3.11 by fixing typing-related bu…

…gs (#1418)

Fixes hidden bugs in `eve.datamodels` and `eve.extended_typing` to support Python 3.11.

Actual bug fixes:
- Previous fix to support `typing.Any` implementation as a class (python/cpython@5a4973e) didn't work in 3.11. 
- Partially concretization of generic datamodels replacing typevars was broken.
- Partially concretization of generic datamodels leaving some parameters as typevars was broken.

Other changes:
- Add python 3.11 as supported version.
- Remove dead code in comments.
- Fix some imports style to comply with our coding guidelines.

.github/workflows/daily-ci.yml

@@ -14,7 +14,7 @@ jobs:
   daily-ci:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         tox-module-factor: ["cartesian", "eve", "next", "storage"]
         os: ["ubuntu-latest"]
         requirements-file: ["requirements-dev.txt", "min-requirements-test.txt", "min-extra-requirements-test.txt"]

.github/workflows/test-cartesian-fallback.yml

@@ -16,7 +16,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         backends: [internal-cpu, dace-cpu]
 
     steps:

.github/workflows/test-cartesian.yml

@@ -23,7 +23,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         backends: [internal-cpu, dace-cpu]
     steps:
     - uses: actions/checkout@v2

.github/workflows/test-eve-fallback.yml

@@ -17,7 +17,7 @@ jobs:
   test-eve:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         os: ["ubuntu-latest"]
 
     runs-on: ${{ matrix.os }}

.github/workflows/test-eve.yml

@@ -20,7 +20,7 @@ jobs:
   test-eve:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         os: ["ubuntu-latest"]
       fail-fast: false
 
@@ -68,4 +68,3 @@ jobs:
     #   with:
     #     name: info-py${{ matrix.python-version }}-${{ matrix.os }}
     #     path: info.txt
-

.github/workflows/test-next-fallback.yml

@@ -15,7 +15,7 @@ jobs:
   test-next:
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.10", "3.11"]
         tox-env-factor: ["nomesh", "atlas"]
         os: ["ubuntu-latest"]
 

.github/workflows/test-next.yml

@@ -18,7 +18,7 @@ jobs:
   test-next:
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.10", "3.11"]
         tox-env-factor: ["nomesh", "atlas"]
         os: ["ubuntu-latest"]
       fail-fast: false

.github/workflows/test-storage-fallback.yml

@@ -18,7 +18,7 @@ jobs:
   test-storage:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         backends: [internal-cpu, dace-cpu]
         os: ["ubuntu-latest"]
 

.github/workflows/test-storage.yml

@@ -21,7 +21,7 @@ jobs:
   test-storage:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         backends: [internal-cpu, dace-cpu]
         os: ["ubuntu-latest"]
       fail-fast: false
@@ -70,4 +70,3 @@ jobs:
     #   with:
     #     name: info-py${{ matrix.python-version }}-${{ matrix.os }}
     #     path: info.txt
-

src/gt4py/eve/datamodels/core.py

@@ -883,17 +883,6 @@ def _substitute_typevars(
         return type_params_map[type_hint], True
     elif getattr(type_hint, "__parameters__", []):
         return type_hint[tuple(type_params_map[tp] for tp in type_hint.__parameters__)], True
-        # TODO(egparedes): WIP fix for partial specialization
-        # # Type hint is a generic model: replace all the concretized type vars
-        # noqa: e800 replaced = False
-        # noqa: e800 new_args = []
-        # noqa: e800 for tp in type_hint.__parameters__:
-        # noqa: e800     if tp in type_params_map:
-        # noqa: e800         new_args.append(type_params_map[tp])
-        # noqa: e800         replaced = True
-        # noqa: e800     else:
-        # noqa: e800         new_args.append(type_params_map[tp])
-        # noqa: e800 return type_hint[tuple(new_args)], replaced
     else:
         return type_hint, False
 
@@ -981,21 +970,14 @@ def __class_getitem__(
         """
         type_args: Tuple[Type] = args if isinstance(args, tuple) else (args,)
         concrete_cls: Type[DataModelT] = concretize(cls, *type_args)
-        res = xtyping.StdGenericAliasType(concrete_cls, type_args)
-        if sys.version_info < (3, 9):
-            # in Python 3.8, xtyping.StdGenericAliasType (aka typing._GenericAlias)
-            # does not copy all required `__dict__` entries, so do it manually
-            for k, v in concrete_cls.__dict__.items():
-                if k not in res.__dict__:
-                    res.__dict__[k] = v
-        return res
+        return concrete_cls
 
     return classmethod(__class_getitem__)
 
 
 def _make_type_converter(type_annotation: TypeAnnotation, name: str) -> TypeConverter[_T]:
-    # TODO(egparedes): if a "typing tree" structure is implemented, refactor this code as a tree traversal.
-    #
+    # TODO(egparedes): if a "typing tree" structure is implemented, refactor this code
+    # as a tree traversal.
     if xtyping.is_actual_type(type_annotation) and not isinstance(None, type_annotation):
         assert not xtyping.get_args(type_annotation)
         assert isinstance(type_annotation, type)
@@ -1316,11 +1298,7 @@ def _make_concrete_with_cache(
     # Replace field definitions with the new actual types for generic fields
     type_params_map = dict(zip(datamodel_cls.__parameters__, type_args))
     model_fields = getattr(datamodel_cls, MODEL_FIELD_DEFINITIONS_ATTR)
-    new_annotations = {
-        # TODO(egparedes): ?
-        # noqa: e800 "__args__": "ClassVar[Tuple[Union[Type, TypeVar], ...]]",
-        # noqa: e800 "__parameters__": "ClassVar[Tuple[TypeVar, ...]]",
-    }
+    new_annotations = {}
 
     new_field_c_attrs = {}
     for field_name, field_type in xtyping.get_type_hints(datamodel_cls).items():
@@ -1353,8 +1331,16 @@ def _make_concrete_with_cache(
         "__module__": module if module else datamodel_cls.__module__,
         **new_field_c_attrs,
     }
-
     concrete_cls = type(class_name, (datamodel_cls,), namespace)
+
+    # Update the tuple of generic parameters in the new class, in case
+    # this is a partial concretization
+    assert hasattr(concrete_cls, "__parameters__")
+    concrete_cls.__parameters__ = tuple(
+        type_params_map[tp_var]
+        for tp_var in datamodel_cls.__parameters__
+        if isinstance(type_params_map[tp_var], typing.TypeVar)
+    )
     assert concrete_cls.__module__ == module or not module
 
     if MODEL_FIELD_DEFINITIONS_ATTR not in concrete_cls.__dict__:

src/gt4py/eve/extended_typing.py

@@ -493,7 +493,7 @@ def _patched_proto_hook(other):  # type: ignore[no-untyped-def]
     if isinstance(_typing.Any, type):  # Python >= 3.11
         _ArtefactTypes = (*_ArtefactTypes, _typing.Any)
 
-# `Any` is a class since typing_extensions >= 4.4
+# `Any` is a class since typing_extensions >= 4.4 and Python 3.11
 if (typing_exts_any := getattr(_typing_extensions, "Any", None)) is not _typing.Any and isinstance(
     typing_exts_any, type
 ):
@@ -504,11 +504,13 @@ def is_actual_type(obj: Any) -> TypeGuard[Type]:
     """Check if an object has an actual type and instead of a typing artefact like ``GenericAlias`` or ``Any``.
 
     This is needed because since Python 3.9:
-        ``isinstance(types.GenericAlias(),  type) is True``
+        ``isinstance(types.GenericAlias(), type) is True``
     and since Python 3.11:
-        ``isinstance(typing.Any,  type) is True``
+        ``isinstance(typing.Any, type) is True``
     """
-    return isinstance(obj, type) and type(obj) not in _ArtefactTypes
+    return (
+        isinstance(obj, type) and (obj not in _ArtefactTypes) and (type(obj) not in _ArtefactTypes)
+    )
 
 
 if hasattr(_typing_extensions, "Any") and _typing.Any is not _typing_extensions.Any:  # type: ignore[attr-defined] # _typing_extensions.Any only from >= 4.4
@@ -641,9 +643,12 @@ def get_partial_type_hints(
             resolved_hints = get_type_hints(  # type: ignore[call-arg]  # Python 3.8 does not define `include-extras`
                 obj, globalns=globalns, localns=localns, include_extras=include_extras
             )
-            hints.update(resolved_hints)
+            hints[name] = resolved_hints[name]
         except NameError as error:
             if isinstance(hint, str):
+                # This conversion could be probably skipped in Python versions containing
+                # the fix applied in bpo-41370. Check:
+                # https://github.com/python/cpython/commit/b465b606049f6f7dd0711cb031fdaa251818741a#diff-ddb987fca5f5df0c9a2f5521ed687919d70bb3d64eaeb8021f98833a2a716887R344
                 hints[name] = ForwardRef(hint)
             elif isinstance(hint, (ForwardRef, _typing.ForwardRef)):
                 hints[name] = hint

tests/eve_tests/unit_tests/test_datamodels.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import enum
+import numbers
 import types
 import typing
 from typing import Set  # noqa: F401 # imported but unused (used in exec() context)
@@ -1150,66 +1151,80 @@ class PartialGenericModel(datamodels.GenericDataModel, Generic[T]):
         with pytest.raises(TypeError, match="'PartialGenericModel__int.value'"):
             PartialGenericModel__int(value=["1"])
 
-    def test_partial_specialization(self):
-        class PartialGenericModel(datamodels.GenericDataModel, Generic[T, U]):
+    def test_partial_concretization(self):
+        class BaseGenericModel(datamodels.GenericDataModel, Generic[T, U]):
             value: List[Tuple[T, U]]
 
-        PartialGenericModel(value=[])
-        PartialGenericModel(value=[("value", 3)])
-        PartialGenericModel(value=[(1, "value")])
-        PartialGenericModel(value=[(-1.0, "value")])
-        with pytest.raises(TypeError, match="'PartialGenericModel.value'"):
-            PartialGenericModel(value=1)
-        with pytest.raises(TypeError, match="'PartialGenericModel.value'"):
-            PartialGenericModel(value=(1, 2))
-        with pytest.raises(TypeError, match="'PartialGenericModel.value'"):
-            PartialGenericModel(value=[()])
-        with pytest.raises(TypeError, match="'PartialGenericModel.value'"):
-            PartialGenericModel(value=[(1,)])
+        assert len(BaseGenericModel.__parameters__) == 2
+
+        BaseGenericModel(value=[])
+        BaseGenericModel(value=[("value", 3)])
+        BaseGenericModel(value=[(1, "value")])
+        BaseGenericModel(value=[(-1.0, "value")])
+        with pytest.raises(TypeError, match="'BaseGenericModel.value'"):
+            BaseGenericModel(value=1)
+        with pytest.raises(TypeError, match="'BaseGenericModel.value'"):
+            BaseGenericModel(value=(1, 2))
+        with pytest.raises(TypeError, match="'BaseGenericModel.value'"):
+            BaseGenericModel(value=[()])
+        with pytest.raises(TypeError, match="'BaseGenericModel.value'"):
+            BaseGenericModel(value=[(1,)])
+
+        PartiallyConcretizedGenericModel = BaseGenericModel[int, U]
+
+        assert len(PartiallyConcretizedGenericModel.__parameters__) == 1
+
+        PartiallyConcretizedGenericModel(value=[])
+        PartiallyConcretizedGenericModel(value=[(1, 2)])
+        PartiallyConcretizedGenericModel(value=[(1, "value")])
+        PartiallyConcretizedGenericModel(value=[(1, (11, 12))])
+        with pytest.raises(TypeError, match=".value'"):
+            PartiallyConcretizedGenericModel(value=1)
+        with pytest.raises(TypeError, match=".value'"):
+            PartiallyConcretizedGenericModel(value=(1, 2))
+        with pytest.raises(TypeError, match=".value'"):
+            PartiallyConcretizedGenericModel(value=[1.0])
+        with pytest.raises(TypeError, match=".value'"):
+            PartiallyConcretizedGenericModel(value=["1"])
 
-        print(f"{PartialGenericModel.__parameters__=}")
-        print(f"{hasattr(PartialGenericModel ,'__args__')=}")
+        FullyConcretizedGenericModel = PartiallyConcretizedGenericModel[str]
 
-        PartiallySpecializedGenericModel = PartialGenericModel[int, U]
-        print(f"{PartiallySpecializedGenericModel.__datamodel_fields__=}")
-        print(f"{PartiallySpecializedGenericModel.__parameters__=}")
-        print(f"{PartiallySpecializedGenericModel.__args__=}")
+        assert len(FullyConcretizedGenericModel.__parameters__) == 0
 
-        PartiallySpecializedGenericModel(value=[])
-        PartiallySpecializedGenericModel(value=[(1, 2)])
-        PartiallySpecializedGenericModel(value=[(1, "value")])
-        PartiallySpecializedGenericModel(value=[(1, (11, 12))])
+        FullyConcretizedGenericModel(value=[])
+        FullyConcretizedGenericModel(value=[(1, "value")])
+        with pytest.raises(TypeError, match=".value'"):
+            FullyConcretizedGenericModel(value=1)
+        with pytest.raises(TypeError, match=".value'"):
+            FullyConcretizedGenericModel(value=(1, 2))
         with pytest.raises(TypeError, match=".value'"):
-            PartiallySpecializedGenericModel(value=1)
+            FullyConcretizedGenericModel(value=[1.0])
         with pytest.raises(TypeError, match=".value'"):
-            PartiallySpecializedGenericModel(value=(1, 2))
+            FullyConcretizedGenericModel(value=["1"])
         with pytest.raises(TypeError, match=".value'"):
-            PartiallySpecializedGenericModel(value=[1.0])
+            FullyConcretizedGenericModel(value=1)
         with pytest.raises(TypeError, match=".value'"):
-            PartiallySpecializedGenericModel(value=["1"])
-
-        # TODO(egparedes): after fixing partial nested datamodel specialization
-        # noqa: e800 FullySpecializedGenericModel = PartiallySpecializedGenericModel[str]
-        # noqa: e800 print(f"{FullySpecializedGenericModel.__datamodel_fields__=}")
-        # noqa: e800 print(f"{FullySpecializedGenericModel.__parameters__=}")
-        # noqa: e800 print(f"{FullySpecializedGenericModel.__args__=}")
-
-        # noqa: e800 FullySpecializedGenericModel(value=[])
-        # noqa: e800 FullySpecializedGenericModel(value=[(1, "value")])
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=1)
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=(1, 2))
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=[1.0])
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=["1"])
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=1)
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=[(1, 2)])
-        # noqa: e800 with pytest.raises(TypeError, match=".value'"):
-        # noqa: e800     FullySpecializedGenericModel(value=[(1, (11, 12))])
+            FullyConcretizedGenericModel(value=[(1, 2)])
+        with pytest.raises(TypeError, match=".value'"):
+            FullyConcretizedGenericModel(value=[(1, (11, 12))])
+
+    def test_partial_concretization_with_typevar(self):
+        class PartialGenericModel(datamodels.GenericDataModel, Generic[T]):
+            a: T
+            values: List[T]
+
+        B = TypeVar("B", bound=numbers.Number)
+        PartiallyConcretizedGenericModel = PartialGenericModel[B]
+
+        PartiallyConcretizedGenericModel(a=1, values=[2, 3])
+        PartiallyConcretizedGenericModel(a=-1.32, values=[2.2, 3j])
+
+        with pytest.raises(TypeError, match=".a'"):
+            PartiallyConcretizedGenericModel(a="1", values=[2, 3])
+        with pytest.raises(TypeError, match=".values'"):
+            PartiallyConcretizedGenericModel(a=1, values=[1, "2"])
+        with pytest.raises(TypeError, match=".values'"):
+            PartiallyConcretizedGenericModel(a=1, values=(1, 2))
 
     # Reuse sample_type_data from test_field_type_hint
     @pytest.mark.parametrize(["type_hint", "valid_values", "wrong_values"], SAMPLE_TYPE_DATA)

tests/eve_tests/unit_tests/test_type_validation.py

@@ -28,6 +28,7 @@
 )
 from gt4py.eve.extended_typing import (
     Any,
+    Callable,
     Dict,
     Final,
     ForwardRef,
@@ -41,8 +42,8 @@
 )
 
 
-VALIDATORS: Final = [type_val.simple_type_validator]
-FACTORIES: Final = [type_val.simple_type_validator_factory]
+VALIDATORS: Final[list[Callable]] = [type_val.simple_type_validator]
+FACTORIES: Final[list[Callable]] = [type_val.simple_type_validator_factory]
 
 
 class SampleEnum(enum.Enum):