Skip to content

Commit

Permalink
introduce DagsterModel (#20637)
Browse files Browse the repository at this point in the history
## Summary & Motivation

Introduces a `StrictModel` class, which is a subclass of
`pydantic.BaseModel` that's frozen and doesn't allow default constructor
arguments that aren't class members.

The idea is for this to eventually replace all our uses of `NamedTuple`,
`BaseModel`, and `@dataclass`.

This is a replacement for the unmerged
#20461.

Downstream PRs that demonstrate its usage:
- #20641
- #20643
- #20638

## How I Tested These Changes
  • Loading branch information
sryza authored and PedramNavid committed Mar 28, 2024
1 parent 4189315 commit 276804c
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 2 deletions.
17 changes: 17 additions & 0 deletions python_modules/dagster/dagster/_model/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from typing import Any

from pydantic import BaseModel


class DagsterModel(BaseModel):
"""Standardizes on Pydantic settings that are stricter than the default.
- Frozen, to avoid complexity caused by mutation.
- extra=forbid, to avoid bugs caused by accidentally constructing with the wrong arguments.
"""

def __init__(self, **data: Any) -> None:
super().__init__(**data)

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._model import DagsterModel
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 SomeDagsterModel(DagsterModel):
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: SomeDagsterModel

o = DataclassObj(
"woo",
4,
InnerDataclass(1.2),
SomeNT([1, 2, 3]),
SomeModel(id=4, name="zuck"),
SomeDagsterModel(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
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
import pytest
from dagster._model import DagsterModel
from pydantic import ValidationError


def test_override_constructor_in_subclass():
class MyClass(DagsterModel):
foo: str
bar: int

def __init__(self, foo: str, bar: int):
super().__init__(foo=foo, bar=bar)

MyClass(foo="fdsjk", bar=4)


def test_override_constructor_in_subclass_different_arg_names():
class MyClass(DagsterModel):
foo: str
bar: int

def __init__(self, fooarg: str, bararg: int):
super().__init__(foo=fooarg, bar=bararg)

MyClass(fooarg="fdsjk", bararg=4)


def test_override_constructor_in_subclass_wrong_type():
class MyClass(DagsterModel):
foo: str
bar: int

def __init__(self, foo: str, bar: str):
super().__init__(foo=foo, bar=bar)

with pytest.raises(ValidationError):
MyClass(foo="fdsjk", bar="fdslk")

0 comments on commit 276804c

Please sign in to comment.