chore: Fix mypy linting errors

Signed-off-by: Mark Wiebe <[email protected]>

Author: mwiebe

src/openjd/model/_convert_pydantic_error.py

@@ -1,6 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-from typing import Type
+from typing import Type, Union
 from pydantic import BaseModel
 from pydantic_core import ErrorDetails
 from inspect import getmodule
@@ -50,7 +50,7 @@ def _error_dict_to_str(root_model: Type[BaseModel], error_details: ErrorDetails)
     return f"{_loc_to_str(root_model, loc)}:\n\t{msg}"
 
 
-def _loc_to_str(root_model: Type[BaseModel], loc: tuple[int | str, ...]) -> str:
+def _loc_to_str(root_model: Type[BaseModel], loc: tuple[Union[int, str], ...]) -> str:
     model_module = getmodule(root_model)
 
     # If a nested error is from a root validator, then just report the error as being

src/openjd/model/_format_strings/_format_string.py

@@ -20,6 +20,7 @@ class ExpressionInfo:
 
 class FormatStringError(ValueError):
     def __init__(self, *, string: str, start: int, end: int, expr: str = "", details: str = ""):
+        self.input = string
         expression = f"Expression: {expr}. " if expr else ""
         reason = f"Reason: {details}." if details else ""
         msg = (

src/openjd/model/_internal/_create_job.py

@@ -1,6 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-from typing import Any, Union
+from typing import cast, Any, Union
 
 from pydantic import ValidationError
 from pydantic_core import InitErrorDetails
@@ -71,14 +71,18 @@ def instantiate_model(  # noqa: C901
         except ValidationError as exc:
             # Convert the ErrorDetails to InitErrorDetails by excluding the 'msg'
             for error_details in exc.errors():
-                errors.append(
-                    InitErrorDetails(
-                        **{key: value for key, value in error_details.items() if key != "msg"}
-                    )
-                )
+                init_error_details = {
+                    key: value for key, value in error_details.items() if key != "msg"
+                }
+                errors.append(cast(InitErrorDetails, init_error_details))
         except FormatStringError as exc:
             errors.append(
-                InitErrorDetails(type="value_error", loc=loc, ctx={"error": ValueError(str(exc))})
+                InitErrorDetails(
+                    type="value_error",
+                    loc=loc,
+                    ctx={"error": ValueError(str(exc))},
+                    input=exc.input,
+                )
             )
 
     if errors:
@@ -105,10 +109,10 @@ def instantiate_model(  # noqa: C901
             init_error_details = {}
             for key, value in error_details.items():
                 if key == "loc":
-                    init_error_details["loc"] = loc + value
+                    init_error_details["loc"] = loc + cast(tuple, value)
                 elif key != "msg":
                     init_error_details[key] = value
-            errors.append(InitErrorDetails(**init_error_details))
+            errors.append(cast(InitErrorDetails, init_error_details))
         raise ValidationError.from_exception_data(
             title=model.__class__.__name__, line_errors=errors
         )
@@ -183,15 +187,17 @@ def _instantiate_list_field(  # noqa: C901
             except ValidationError as exc:
                 # Convert the ErrorDetails to InitErrorDetails by excluding the 'msg'
                 for error_details in exc.errors():
-                    errors.append(
-                        InitErrorDetails(
-                            **{key: value for key, value in error_details.items() if key != "msg"}
-                        )
-                    )
+                    init_error_details = {
+                        key: value for key, value in error_details.items() if key != "msg"
+                    }
+                    errors.append(cast(InitErrorDetails, init_error_details))
             except FormatStringError as exc:
                 errors.append(
                     InitErrorDetails(
-                        type="value_error", loc=loc, ctx={"error": ValueError(str(exc))}
+                        type="value_error",
+                        loc=loc,
+                        ctx={"error": ValueError(str(exc))},
+                        input=exc.input,
                     )
                 )
     else:
@@ -211,15 +217,17 @@ def _instantiate_list_field(  # noqa: C901
             except ValidationError as exc:
                 # Convert the ErrorDetails to InitErrorDetails by excluding the 'msg'
                 for error_details in exc.errors():
-                    errors.append(
-                        InitErrorDetails(
-                            **{key: value for key, value in error_details.items() if key != "msg"}
-                        )
-                    )
+                    init_error_details = {
+                        key: value for key, value in error_details.items() if key != "msg"
+                    }
+                    errors.append(cast(InitErrorDetails, init_error_details))
             except FormatStringError as exc:
                 errors.append(
                     InitErrorDetails(
-                        type="value_error", loc=loc, ctx={"error": ValueError(str(exc))}
+                        type="value_error",
+                        loc=loc,
+                        ctx={"error": ValueError(str(exc))},
+                        input=exc.input,
                     )
                 )
 
@@ -266,14 +274,18 @@ def _instantiate_dict_field(
         except ValidationError as exc:
             # Convert the ErrorDetails to InitErrorDetails by excluding the 'msg'
             for error_details in exc.errors():
-                errors.append(
-                    InitErrorDetails(
-                        **{key: value for key, value in error_details.items() if key != "msg"}
-                    )
-                )
+                init_error_details = {
+                    key: value for key, value in error_details.items() if key != "msg"
+                }
+                errors.append(cast(InitErrorDetails, init_error_details))
         except FormatStringError as exc:
             errors.append(
-                InitErrorDetails(type="value_error", loc=loc, ctx={"error": ValueError(str(exc))})
+                InitErrorDetails(
+                    type="value_error",
+                    loc=loc,
+                    ctx={"error": ValueError(str(exc))},
+                    input=exc.input,
+                )
             )
 
     if errors:

src/openjd/model/_internal/_variable_reference_validation.py

@@ -2,11 +2,12 @@
 
 from collections import defaultdict
 import typing
-from typing import Any, Optional, Type, Literal, Union
+from typing import cast, Any, Optional, Type, Literal, Union
 from inspect import isclass
 
+from pydantic import Discriminator
 from pydantic_core import InitErrorDetails
-from pydantic.fields import FieldInfo
+from pydantic.fields import FieldInfo, ModelPrivateAttr
 
 from .._types import OpenJDModel, ResolutionScope
 from .._format_strings import FormatString, FormatStringError
@@ -101,36 +102,18 @@
 #    of their fields variable definitions are passed in to.
 #    - Information on this is encoded in the model's `_template_variable_sources` field. See the comment for this field in the
 #      OpenJDModel base class for information on this property
-
-
-# TODO: Rewrite this comment for Pydantic 2
 # 4. Since this validation is a pre-validator, we basically have to re-implement a fragment of Pydantic's model parser for this
-#    depth first traversal. Thus, you'll need to know the following about Pydantic v1.x's data model and parser to understand this
+#    depth first traversal. Thus, you'll need to know the following about Pydantic v2's data model and parser to understand this
 #    implementation:
 #    a) All models are derived from pydantic.BaseModel
 #    b) pydantic.BaseModel.model_fields: dict[str, FieldInfo] is injected into all BaseModels by pydantic's BaseModel metaclass.
 #       This member is what gives pydantic information about each of the fields defined in the model class. The key of the dict is the
 #       name of the field in the model.
 #    c) pydantic.FieldInfo describes the type information about a model's field:
-#        i) pydantic.FieldInfo.annotation is the type annotating the field
-#                Literal[...] means that it's a singleton type.
-#                list[...] means that it's a list type.
-#                dict[...] means that it's a dict type.
-#                etc.
-#       ii) pydantic.FieldInfo.annotation gives you the type of the field; this is only useful for scalar singleton fields.
-#       iii) pydantic.FieldInfo.sub_fields: Optional[list[pydantic.FieldInfo]] exists for list, dictionary, and union-typed singleton
-#            fields:
-#            1. For SHAPE_LIST: sub_fields has length 1, and its element is the FieldInfo for the elements of the list.
-#            2. For SHAPE_DICT: sub_fields has length 1, and its element is the FieldInfo for the value-type of the dict.
-#            3. For SHAPE_SINGLETON:
-#                 a) For scalar-typed fields: sub_fields is None
-#                 b) For union-typed fields: sub_fields is a list of all of the types in the union
-#       iv) For discriminated unions:
-#            1. pydantic.FieldInfo.discriminator_key: Optional[str] exists and it gives the name of the submodel field used to
-#               determine which type of the union a given data value is.
-#            2. pydantic.sub_fields_mapping: Optional[dict[str,pydantic.FieldInfo]] exists and can be used to find the unioned type
-#               for a given discriminator value.
-#
+#       i) pydantic.FieldInfo.annotation gives you the type of the field; The structure of the field is contained
+#          in this type, including typing.Annotated values for discriminated unions. Both the definition collection and
+#          validation recursively unwraps these types along with the values. The pydantic.FieldInfo also includes a discriminator
+#          value, so the code handles both cases.
 
 
 class ScopedSymtabs(defaultdict):
@@ -198,7 +181,7 @@ def _validate_model_template_variable_references(
     symbol_prefix: str,
     symbols: ScopedSymtabs,
     loc: tuple,
-    discriminator: Optional[str] = None,
+    discriminator: Union[str, Discriminator, None] = None,
 ) -> list[InitErrorDetails]:
     """Inner implementation of prevalidate_model_template_variable_references().
 
@@ -290,14 +273,18 @@ def _validate_model_template_variable_references(
 
     # Unwrap a discriminated union to the selected type
     if model_origin is Union and discriminator is not None:
-        return _validate_model_template_variable_references(
-            _get_model_for_singleton_value(model, value, discriminator),
-            value,
-            current_scope,
-            symbol_prefix,
-            symbols,
-            loc,
-        )
+        unioned_model = _get_model_for_singleton_value(model, value, discriminator)
+        if unioned_model is not None:
+            return _validate_model_template_variable_references(
+                unioned_model,
+                value,
+                current_scope,
+                symbol_prefix,
+                symbols,
+                loc,
+            )
+        else:
+            return []
 
     if isclass(model) and issubclass(model, FormatString):
         if isinstance(value, str):
@@ -310,8 +297,9 @@ def _validate_model_template_variable_references(
 
     # Does this cls change the variable reference scope for itself and its children? If so, then update
     # our scope.
-    if model._template_variable_scope.get_default() is not None:
-        current_scope = model._template_variable_scope.get_default()
+    model_override_scope = cast(ModelPrivateAttr, model._template_variable_scope).get_default()
+    if model_override_scope is not None:
+        current_scope = model_override_scope
 
     # Apply any changes that this node makes to the template variable prefix.
     #  e.g. It may change "Env." to "Env.File."
@@ -331,9 +319,10 @@ def _validate_model_template_variable_references(
     # Recursively validate the contents of FormatStrings within the model.
     for field_name, field_info in model.model_fields.items():
         field_value = value.get(field_name)
-        if field_value is None:
+        field_model = field_info.annotation
+        if field_value is None or field_model is None:
             continue
-        if typing.get_origin(field_info.annotation) is Literal:
+        if typing.get_origin(field_model) is Literal:
             # Literals aren't format strings and cannot be recursed in to; skip them.
             continue
 
@@ -350,7 +339,7 @@ def _validate_model_template_variable_references(
 
         errors.extend(
             _validate_model_template_variable_references(
-                field_info.annotation,
+                field_model,
                 field_value,
                 current_scope,
                 symbol_prefix,
@@ -381,13 +370,15 @@ def _check_format_string(
             try:
                 expr.expression.validate_symbol_refs(symbols=scoped_symbols)
             except ValueError as exc:
-                errors.append(InitErrorDetails(type="value_error", loc=loc, ctx={"error": exc}))
+                errors.append(
+                    InitErrorDetails(type="value_error", loc=loc, ctx={"error": exc}, input=value)
+                )
     return errors
 
 
 def _get_model_for_singleton_value(
-    model: Any, value: Any, discriminator: Optional[str] = None
-) -> Optional[FieldInfo]:
+    model: Any, value: Any, discriminator: Union[str, Discriminator, None] = None
+) -> Optional[Type]:
     """Given a FieldInfo and the value that we're given for that field, determine
     the actual Model for the value in the event that the FieldInfo may be for
     a discriminated union."""
@@ -453,7 +444,7 @@ def _collect_variable_definitions(  # noqa: C901  (suppress: too complex)
     current_scope: ResolutionScope,
     symbol_prefix: str,
     recursive_pruning: bool = True,
-    discriminator: Optional[str] = None,
+    discriminator: Union[str, Discriminator, None] = None,
 ) -> dict[str, ScopedSymtabs]:
     """Collects the names of variables that each field of this model object provides.
 
@@ -518,12 +509,16 @@ def _collect_variable_definitions(  # noqa: C901  (suppress: too complex)
 
     # Unwrap a discriminated union to the selected type
     if model_origin is Union and discriminator is not None:
-        return _collect_variable_definitions(
-            _get_model_for_singleton_value(model, value, discriminator),
-            value,
-            current_scope,
-            symbol_prefix,
-        )
+        unioned_model = _get_model_for_singleton_value(model, value, discriminator)
+        if unioned_model is not None:
+            return _collect_variable_definitions(
+                unioned_model,
+                value,
+                current_scope,
+                symbol_prefix,
+            )
+        else:
+            return {"__export__": ScopedSymtabs()}
 
     # Anything except for an OpenJDModel returns an empty result
     if not isclass(model) or not issubclass(model, OpenJDModel):
@@ -577,10 +572,10 @@ def _collect_variable_definitions(  # noqa: C901  (suppress: too complex)
     # Collect the variable definitions exported by the fields of the model
     for field_name, field_info in model.model_fields.items():
         field_value = value.get(field_name)
-        if field_value is None:
+        field_model = field_info.annotation
+        if field_value is None or field_model is None:
             continue
 
-        field_model = field_info.annotation
         discriminator = field_info.discriminator
 
         symbols[field_name] = _collect_variable_definitions(

src/openjd/model/_parse.py

@@ -9,7 +9,7 @@
 import yaml
 from pydantic import BaseModel
 from pydantic import ValidationError as PydanticValidationError
-from pydantic_core import ErrorDetails
+from pydantic_core import ErrorDetails, InitErrorDetails
 
 from ._errors import DecodeValidationError
 from ._types import EnvironmentTemplate, JobTemplate, OpenJDModel, TemplateSpecificationVersion
@@ -58,8 +58,14 @@ def _parse_model(*, model: Type[T], obj: Any) -> T:
             result = cast(T, cast(BaseModel, model).model_validate(obj))
         except PydanticValidationError as exc:
             if prevalidator_error is not None:
+                errors = list[InitErrorDetails]()
+                for error_details in exc.errors() + prevalidator_error.errors():
+                    init_error_details = {
+                        key: value for key, value in error_details.items() if key != "msg"
+                    }
+                    errors.append(cast(InitErrorDetails, init_error_details))
                 raise PydanticValidationError.from_exception_data(
-                    exc.title, exc.errors() + prevalidator_error.errors()
+                    title=exc.title, line_errors=errors
                 )
             else:
                 raise