Repository: omni-us/jsonargparse

Issue:
Dot syntax not working correctly for pydantic 2 models
<!-- If you like this project, please ⭐ star it https://github.com/omni-us/jsonargparse/ -->

## 🐛 Bug report

<!-- A clear and concise description of the bug. -->

### To reproduce

```py
from jsonargparse import CLI
from pydantic import BaseModel, Field
from typing import Annotated, Literal, Union


class PingTask(BaseModel):
    type: Literal["ping"] = "ping"


class PongTask(BaseModel):
    type: Literal["pong"] = "pong"

Task = Annotated[
    Union[PingTask, PongTask],
    Field(discriminator="type"),
]

def train_func(model: Task) -> None:
    print(model)


if __name__ == '__main__':
    CLI(train_func, as_positional=False)
```

It works well when the model argument is specified as a json string:

```bash
$ python x.py --model '{"type": "ping"}'
type='ping'
```

However, it doesn't work using the dot syntax:

```bash
$ python x.py --model.type ping                                                                                                

usage: x.py [-h] [--config CONFIG] [--print_config[=flags]] --model MODEL
error: Parser key "model":
  1 validation error for tagged-union[PingTask,PongTask]
    Unable to extract tag using discriminator 'type' [type=union_tag_not_found, input_value=NestedArg(key='type', val='ping'), input_type=NestedArg]
      For further information visit https://errors.pydantic.dev/2.9/v/union_tag_not_found. Got value: NestedArg(key='type', val='ping')
```

### Expected behavior

<!-- Describe how the behavior or output of the reproduction snippet should be. -->

### Environment

<!-- Fill in the list below. -->

- jsonargparse version: 4.33.2
- Python version: 3.11.9
- How jsonargparse was installed: pip install jsonargparse
- OS: mac

Discussion:
@mauvilsa: Thank you for reporting! I took a quick look. Seems that this happens with pydantic 2 but works fine with pydantic 1. Not sure why. Will look in more detail later.@mauvilsa: An update on this. What is not supported is this kind of `Annotated` type. Not sure yet what to do to support that. However, in the example above it is not clear why such a type is necessary. At least from the perspective of jsonargparse, the dot syntax would work fine if you do without the `Annotated` and instead:

```python
def train_func(model: Union[PingTask, PongTask]) -> None:
    print(model)
```

The `type` discriminator is already implied by the literals. So there is no need for the `Annotated`.@mauvilsa: There might be other cases in which a discriminated union might make sense, unlike the one above. If you still think that there is a need to support it, please provide a new example that motivates it.@findmyway: > The type discriminator is already implied by the literals. So there is no need for the Annotated.


Hi @mauvilsa , thanks for the update.


The `Annotated` is used to correctly deserialize from the raw JSON string. For example:

```py
s = '{"type": "ping"}'

x = Task.parse_raw(s)
```@mauvilsa: I am not sure what you mean. For me running `Task.parse_raw('{"type": "ping"}')` produces `AttributeError: '_SpecialForm' object has no attribute 'parse_raw'`. And even without running it, `typing.Annotated` is not part of pydantic and it is not for parsing. Why would it have a `parse_raw` method or know how to parse into a pydantic object?@findmyway: Sorry I was hurry to go out and didn't verify the code above. But this blog may explain the reason: https://blog.devgenius.io/deserialize-child-classes-with-pydantic-that-gonna-work-784230e1cf83

> ... So there is no need for the `Annotated`.

I think now I get what you mean above. Indeed, for `jsonargparse` to correctly understand the arg type, we need `Union` type only. But in the real world case, the arg has to be of type `Annotated` with a correct dispatch method.

So I think the best approach here is to add a wrapper and only need to tell `jsonargparse` the basic `Union` type info.

(I'm closing this issue for now because this sounds like a good to have feature)

@mauvilsa: Even after looking at the blog post and the pydantic docs, I still think the `Annotated` is unnecessary. I just merged pull request #667 which for this case simply ignores the `Annotated` and takes only the `Union`, so the example above will now work fine.@findmyway: Thanks!

Solution patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 19c2a294..de173b2b 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -25,6 +25,8 @@ Fixed
 ^^^^^
 - Incorrect instantiation order when instantiation targets share a parent (`#662
   <https://github.com/omni-us/jsonargparse/pull/662>`__).
+- Pydantic discriminated unions handled incorrectly (`#667
+  <https://github.com/omni-us/jsonargparse/pull/667>`__).
 
 
 v4.36.0 (2025-01-17)
diff --git a/jsonargparse/_optionals.py b/jsonargparse/_optionals.py
index 6a02143a..97ed2c8e 100644
--- a/jsonargparse/_optionals.py
+++ b/jsonargparse/_optionals.py
@@ -5,7 +5,7 @@
 from contextlib import contextmanager
 from importlib.metadata import version
 from importlib.util import find_spec
-from typing import Optional
+from typing import Optional, Union
 
 __all__ = [
     "get_config_read_mode",
@@ -348,10 +348,13 @@ def get_module(value):
 
 
 def is_annotated_validator(typehint: type) -> bool:
+    from ._util import get_typehint_origin
+
     return (
         pydantic_support > 1
         and is_annotated(typehint)
         and any(get_module(m) in {"pydantic", "annotated_types"} for m in typehint.__metadata__)  # type: ignore[attr-defined]
+        and get_typehint_origin(typehint.__origin__) != Union  # type: ignore[attr-defined]
     )
 
 


Test patch:
diff --git a/jsonargparse_tests/test_dataclass_like.py b/jsonargparse_tests/test_dataclass_like.py
index 51a990aa..1148f606 100644
--- a/jsonargparse_tests/test_dataclass_like.py
+++ b/jsonargparse_tests/test_dataclass_like.py
@@ -2,7 +2,7 @@
 
 import dataclasses
 import json
-from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, Generic, List, Literal, Optional, Tuple, TypeVar, Union
 from unittest.mock import patch
 
 import pytest
@@ -713,6 +713,18 @@ def test_pydantic_annotated_nested_annotated_dataclass_with_default_factory(pars
         cfg = parser.parse_args(["--n", "{}"])
         assert cfg.n == Namespace(a1=Namespace(a2=1))
 
+    class PingTask(pydantic.BaseModel):
+        type: Literal["ping"] = "ping"
+        attr: str = ""
+
+    class PongTask(pydantic.BaseModel):
+        type: Literal["pong"] = "pong"
+
+    PingPongTask = annotated[
+        Union[PingTask, PongTask],
+        pydantic.Field(discriminator="type"),
+    ]
+
 
 length = "length"
 if pydantic_support:
@@ -806,6 +818,8 @@ def test_subclass(self, parser):
         parser.add_argument("--model", type=PydanticSubModel, default=PydanticSubModel(p1="a"))
         cfg = parser.parse_args(["--model.p3=0.2"])
         assert Namespace(p1="a", p2=3, p3=0.2) == cfg.model
+        init = parser.instantiate_classes(cfg)
+        assert isinstance(init.model, PydanticSubModel)
 
     def test_field_default_factory(self, parser):
         parser.add_argument("--model", type=PydanticFieldFactory)
@@ -831,6 +845,18 @@ def test_annotated_field(self, parser):
             parser.parse_args(["--model.p1=0"])
         ctx.match("model.p1")
 
+    @pytest.mark.skipif(not (annotated and pydantic_support > 1), reason="Annotated is required")
+    def test_field_union_discriminator_dot_syntax(self, parser):
+        parser.add_argument("--model", type=PingPongTask)
+        cfg = parser.parse_args(["--model.type=pong"])
+        assert cfg.model == Namespace(type="pong")
+        init = parser.instantiate_classes(cfg)
+        assert isinstance(init.model, PongTask)
+        cfg = parser.parse_args(["--model.type=ping", "--model.attr=abc"])
+        assert cfg.model == Namespace(type="ping", attr="abc")
+        init = parser.instantiate_classes(cfg)
+        assert isinstance(init.model, PingTask)
+
     @pytest.mark.parametrize(
         ["valid_value", "invalid_value", "cast", "type_str"],
         [