Skip to content

Commit

Permalink
introduce StrictModel
Browse files Browse the repository at this point in the history
branch-name: strict-model
  • Loading branch information
sryza committed Mar 21, 2024
1 parent b15b663 commit 82da062
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 2 deletions.
7 changes: 7 additions & 0 deletions python_modules/dagster/dagster/_core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from weakref import WeakSet

import toposort as toposort_
from pydantic import BaseModel
from typing_extensions import Final

import dagster._check as check
Expand Down Expand Up @@ -196,3 +197,9 @@ def submit(self, fn, *args, **kwargs):
def is_valid_email(email: str) -> bool:
regex = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,7}\b"
return bool(re.fullmatch(regex, email))


class StrictModel(BaseModel):
class Config:
extra = "forbid"
frozen = True
18 changes: 16 additions & 2 deletions python_modules/dagster/dagster_tests/general_tests/test_serdes.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import pydantic
import pytest
from dagster._check import ParameterCheckError, inst_param, set_param
from dagster._core.utils import StrictModel
from dagster._serdes.errors import DeserializationError, SerdesUsageError, SerializationError
from dagster._serdes.serdes import (
EnumSerializer,
Expand Down Expand Up @@ -822,6 +823,11 @@ class SomeModel(pydantic.BaseModel):
id: int
name: str

@_whitelist_for_serdes(test_env)
class SomeStrictModel(StrictModel):
id: int
name: str

@_whitelist_for_serdes(test_env)
@pydantic.dataclasses.dataclass
class DataclassObj:
Expand All @@ -830,8 +836,16 @@ class DataclassObj:
d: InnerDataclass
nt: SomeNT
m: SomeModel

o = DataclassObj("woo", 4, InnerDataclass(1.2), SomeNT([1, 2, 3]), SomeModel(id=4, name="zuck"))
st: SomeStrictModel

o = DataclassObj(
"woo",
4,
InnerDataclass(1.2),
SomeNT([1, 2, 3]),
SomeModel(id=4, name="zuck"),
SomeStrictModel(id=4, name="zuck"),
)
ser_o = serialize_value(o, whitelist_map=test_env)
assert deserialize_value(ser_o, whitelist_map=test_env) == o

Expand Down

0 comments on commit 82da062

Please sign in to comment.