Skip to content

Commit c79e0e4

Browse files
committed
Mixing CoreModel variants in diff_models()
Allow mixing `CoreModel.__request__` and `CoreModel.__response__` in the `diff_models()` utility function. Before: ```python >>> diff_models(M.__request__(a=1), M.__response__(a=2)) Traceback (most recent call last): File "<stdin>", line 1, in <module> File "/src/dstack/_internal/core/services/diff.py", line 37, in diff_models raise TypeError("Both instances must be of the same Pydantic model class.") TypeError: Both instances must be of the same Pydantic model class. ``` After: ```python >>> diff_models(M.__request__(a=1), M.__response__(a=2)) {'a': ModelFieldDiffRequest(old=1, new=2)} ```
1 parent 4c58f5c commit c79e0e4

2 files changed

Lines changed: 131 additions & 2 deletions

File tree

src/dstack/_internal/core/services/diff.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,14 @@ def diff_models(
3333
A dict of changed fields in the form of
3434
`{<field_name>: {"old": old_value, "new": new_value}}`
3535
"""
36-
if type(old) is not type(new):
36+
if not (
37+
type(old) is type(new)
38+
or (
39+
isinstance(old, CoreModel)
40+
and isinstance(new, CoreModel)
41+
and type(old).__response__ is type(new).__response__
42+
)
43+
):
3744
raise TypeError("Both instances must be of the same Pydantic model class.")
3845

3946
if reset is not None:

src/tests/_internal/core/services/test_diff.py

Lines changed: 123 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,128 @@
11
import pytest
2+
from pydantic import BaseModel
23

3-
from dstack._internal.core.services.diff import ModelDiff, ModelFieldDiff, flatten_diff_fields
4+
from dstack._internal.core.models.common import CoreModel
5+
from dstack._internal.core.services.diff import (
6+
ModelDiff,
7+
ModelFieldDiff,
8+
diff_models,
9+
flatten_diff_fields,
10+
)
11+
12+
13+
class TestDiffModels:
14+
class _BaseModelA(BaseModel):
15+
a: int
16+
b: str
17+
18+
class _BaseModelB(BaseModel):
19+
c: int
20+
21+
class _BaseModelAB(_BaseModelA, _BaseModelB):
22+
pass
23+
24+
class _CoreModelA(CoreModel):
25+
a: int
26+
b: str
27+
28+
class _CoreModelB(CoreModel):
29+
c: int
30+
31+
class _CoreModelAB(_CoreModelA, _CoreModelB):
32+
pass
33+
34+
@pytest.mark.parametrize(
35+
("old", "new", "expected"),
36+
[
37+
pytest.param(
38+
_BaseModelA(a=1, b="x"),
39+
_BaseModelA(a=1, b="y"),
40+
{"b": ModelFieldDiff(old="x", new="y")},
41+
id="base-model",
42+
),
43+
pytest.param(
44+
_CoreModelA(a=1, b="x"),
45+
_CoreModelA(a=1, b="y"),
46+
{"b": ModelFieldDiff(old="x", new="y")},
47+
id="core-model",
48+
),
49+
pytest.param(
50+
_BaseModelA(a=1, b="x"),
51+
_BaseModelA(a=1, b="x"),
52+
{},
53+
id="base-model-no-diff",
54+
),
55+
pytest.param(
56+
_CoreModelA(a=1, b="x"),
57+
_CoreModelA(a=1, b="x"),
58+
{},
59+
id="core-model-no-diff",
60+
),
61+
pytest.param(
62+
_CoreModelA.__request__(a=1, b="x"),
63+
_CoreModelA.__request__(a=1, b="y"),
64+
{"b": ModelFieldDiff(old="x", new="y")},
65+
id="core-model-request",
66+
),
67+
pytest.param(
68+
_CoreModelA.__response__(a=1, b="x"),
69+
_CoreModelA.__response__(a=1, b="y"),
70+
{"b": ModelFieldDiff(old="x", new="y")},
71+
id="core-model-response",
72+
),
73+
pytest.param(
74+
_CoreModelA.__request__(a=1, b="x"),
75+
_CoreModelA.__response__(a=1, b="y"),
76+
{"b": ModelFieldDiff(old="x", new="y")},
77+
id="core-model-request-response",
78+
),
79+
pytest.param(
80+
_CoreModelA(a=1, b="x"),
81+
_CoreModelA.__response__(a=1, b="y"),
82+
{"b": ModelFieldDiff(old="x", new="y")},
83+
id="core-model-base-request",
84+
),
85+
pytest.param(
86+
_CoreModelA(a=1, b="x"),
87+
_CoreModelA.__response__(a=1, b="y"),
88+
{"b": ModelFieldDiff(old="x", new="y")},
89+
id="core-model-base-response",
90+
),
91+
],
92+
)
93+
def test_diff_models(self, old: BaseModel, new: BaseModel, expected: ModelDiff) -> None:
94+
assert diff_models(old, new) == expected
95+
96+
@pytest.mark.parametrize(
97+
("old", "new"),
98+
[
99+
pytest.param(
100+
_BaseModelA(a=1, b="x"),
101+
_BaseModelB(c=2),
102+
id="different-base-models",
103+
),
104+
pytest.param(
105+
_BaseModelA(a=1, b="x"),
106+
_BaseModelAB(a=1, b="x", c=2),
107+
id="base-model-and-subclass",
108+
),
109+
pytest.param(
110+
_CoreModelA(a=1, b="x"),
111+
_CoreModelB(c=2),
112+
id="different-core-models",
113+
),
114+
pytest.param(
115+
_CoreModelA(a=1, b="x"),
116+
_CoreModelAB(a=1, b="x", c=2),
117+
id="core-model-and-subclass",
118+
),
119+
],
120+
)
121+
def test_type_mismatch(self, old: BaseModel, new: BaseModel) -> None:
122+
with pytest.raises(
123+
TypeError, match="Both instances must be of the same Pydantic model class."
124+
):
125+
diff_models(old, new)
4126

5127

6128
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)