Force type checker to use Pydantic's V1 API instead (#1432)

## Pull Request Checklist

- [x] Fixes #1435
- [x] Tests added
- [X] Documentation/examples added
- [X] [Good commit messages](https://cbea.ms/git-commit/) and/or PR
title

This is a bit tricky, so I'm not 100% sure if I got this right.

This is the my current understanding of the situation, as described in
this previous [bug
report](#984 (comment)).

- Hera supports pretty much any version of Pydantic from `>=1.7,<3.0`.
- Pydantic V1 and V2 has different APIs and Hera handles this
differently at runtime fine.
- For static type checkers, we need to pick one at type check time using
`typing.TYPE_CHECKING`. This is necessary for both external type
checkers like pyright and editor tooling.
- The current type checker defaults to V2 APIs, even though Hera models
use V1 under the hood.
- This leads to wrong types and editor hints for methods like
`wf.dict()`.

For example, pylance would report this warning in vscode, but it's
actually wrong.

```
The method "dict" in class "BaseModel" is deprecated
  The `dict` method is deprecated; use `model_dump` instead.
```

Because Hera models use v1 API, `dict()` is not deprecated. If the user
switch to `model_dump` as prescribed by the error, this would fail at
runtime since `model_dump` doesn't actually exist on v1 models.

I think the simplest solution here is to switch to v1 types explicitly?
Since this is only used for Hera's base classes, there shouldn't be any
impact elsewhere?

## Test Plan

I can think of a few different things to test with both v1 and v2
version of pydantic.



- [X] Make sure type checks pass, which I guess is covered by this CI.
- [X] The [Original
PR](https://github.com/argoproj-labs/hera/pull/950/files) added these
conditional imports to make sure vscode can show inline hints for the
constructors like WorkflowTemplate. Make sure this is OK in vscode.
- [X] Make sure `dict()` doesn't show a warning.

- ✅ I validated these 3 with my work project that's using Pydantic V2.
- ✅ I validated v1 behaviour with this [tiny project in a
gist](https://gist.github.com/jaseemabid/63e26ba25cd3e346d57ca9df791e3565).

Attaching a vscode screenshot as well.

<img width="772" alt="image"
src="https://github.com/user-attachments/assets/955baacc-ce63-41fd-9339-b81e471284ee"
/>

---------

Signed-off-by: Jaseem Abid <[email protected]>
Co-authored-by: Elliot Gunton <[email protected]>

Author: jaseemabid

src/hera/shared/_pydantic.py

@@ -16,35 +16,41 @@
 # users across both versions.
 
 if _PYDANTIC_VERSION == 2:
-    from pydantic.v1 import (  # type: ignore
+    from pydantic.v1 import (
+        BaseModel as PydanticBaseModel,
         Field,
         PrivateAttr,
         ValidationError,
         root_validator,
         validator,
     )
+    from pydantic.v1.fields import FieldInfo
 else:
     from pydantic import (  # type: ignore[assignment,no-redef]
+        BaseModel as PydanticBaseModel,
         Field,
         PrivateAttr,
         ValidationError,
         root_validator,
         validator,
     )
-
-# TYPE_CHECKING-guarding specifically the `BaseModel` import helps the type checkers
-# provide proper type checking to models. Without this, both mypy and pyright lose
-# native pydantic hinting for `__init__` arguments.
+    from pydantic.fields import FieldInfo  # type: ignore[assignment,no-redef]
+
+
+# TYPE_CHECKING-guarding specifically the `BaseModel` import helps the type
+# checkers provide proper type checking to models. Without this, both mypy and
+# pyright lose native pydantic hinting for `__init__` arguments.
+#
+# Now there can be 2 versions of `BaseModel` at runtime, but here we must choose
+# one for the type checker, otherwise external type checkers as well as editors
+# won't be able to do much. vscode for example would always give 2
+# options when jumping to source instead of just picking one. Pydantic/Pylance
+# wouldn't be able to do much.
+#
+# Since Hera's models are still using v1 API, that's the reasonable choice.
 if TYPE_CHECKING:
-    from pydantic import BaseModel as PydanticBaseModel
-    from pydantic.fields import FieldInfo
-else:
-    if _PYDANTIC_VERSION == 2:
-        from pydantic.v1 import BaseModel as PydanticBaseModel  # type: ignore
-        from pydantic.v1.fields import FieldInfo
-    else:
-        from pydantic import BaseModel as PydanticBaseModel  # type: ignore[assignment,no-redef]
-        from pydantic.fields import FieldInfo
+    from pydantic.v1 import BaseModel as PydanticBaseModel
+    from pydantic.v1.fields import FieldInfo
 
 
 def get_fields(cls: Type[PydanticBaseModel]) -> Dict[str, FieldInfo]:
@@ -55,17 +61,6 @@ def get_fields(cls: Type[PydanticBaseModel]) -> Dict[str, FieldInfo]:
         return cls.__fields__  # type: ignore
 
 
-__all__ = [
-    "BaseModel",
-    "Field",
-    "PrivateAttr",
-    "PydanticBaseModel",  # Export for serialization.py to cover user-defined models
-    "ValidationError",
-    "root_validator",
-    "validator",
-]
-
-
 def get_field_annotations(cls: Type[PydanticBaseModel]) -> Dict[str, Any]:
     return {k: v for k, v in ChainMap(*(get_annotations(c) for c in cls.__mro__)).items()}
 

src/hera/workflows/io/_io_mixins.py

@@ -33,7 +33,7 @@
 if TYPE_CHECKING:
     # We add BaseModel as a parent class of the mixins only when type checking which allows it
     # to be used with either a V1 BaseModel or a V2 BaseModel
-    from pydantic import BaseModel
+    from hera.shared._pydantic import PydanticBaseModel as BaseModel
 else:
     # Subclassing `object` when using the real code (i.e. not type-checking) is basically a no-op
     BaseModel = object  # type: ignore

src/hera/workflows/io/v2.py

@@ -11,7 +11,7 @@
 if _PYDANTIC_VERSION == 2:
     from pydantic import BaseModel
 
-    class Input(InputMixin, BaseModel):
+    class Input(InputMixin, BaseModel):  # type: ignore[misc]
         """Input model usable by the Hera Runner.
 
         Input is a Pydantic model which users can create a subclass of. When a subclass
@@ -20,7 +20,7 @@ class Input(InputMixin, BaseModel):
         for the script_pydantic_io experimental feature.
         """
 
-    class Output(OutputMixin, BaseModel):
+    class Output(OutputMixin, BaseModel):  # type: ignore[misc]
         """Output model usable by the Hera Runner.
 
         Output is a Pydantic model which users can create a subclass of. When a subclass

tests/test_unit/test_pydantic_io_mixins.py

@@ -151,6 +151,24 @@ class Foo(Input):
     )
 
 
+def test_get_as_arguments_unannotated_v1():
+    from hera.workflows.io.v1 import Input as InputV1
+
+    class Foo(InputV1):
+        foo: int
+        bar: str = "a default"
+
+    foo = Foo(foo=1)
+    parameters = foo._get_as_arguments()
+
+    assert parameters == ModelArguments(
+        parameters=[
+            ModelParameter(name="foo", value=1),
+            ModelParameter(name="bar", value="a default"),
+        ],
+    )
+
+
 def test_get_as_arguments_with_pydantic_annotations():
     class Foo(Input):
         foo: Annotated[int, Field(gt=0)]