Skip to content

Commit

Permalink
Do not use the mixin docstring for the schema (#12675)
Browse files Browse the repository at this point in the history
  • Loading branch information
abrookins authored Apr 11, 2024
1 parent fb159b8 commit ef09f72
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 0 deletions.
17 changes: 17 additions & 0 deletions src/prefect/_internal/compatibility/deprecated.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,10 @@

if HAS_PYDANTIC_V2:
from pydantic.v1 import BaseModel, Field, root_validator
from pydantic.v1.schema import default_ref_template
else:
from pydantic import BaseModel, Field, root_validator
from pydantic.schema import default_ref_template

from prefect.utilities.callables import get_call_parameters
from prefect.utilities.importtools import (
Expand Down Expand Up @@ -346,6 +348,21 @@ def dict(self, **kwargs) -> Dict[str, Any]:

return super().dict(**kwargs)

@classmethod
def schema(
cls, by_alias: bool = True, ref_template: str = default_ref_template
) -> Dict[str, Any]:
"""
Don't use the mixin docstring as the description if this class is missing a
docstring.
"""
schema = super().schema(by_alias=by_alias, ref_template=ref_template)

if not cls.__doc__:
schema.pop("description", None)

return schema


def handle_deprecated_infra_overrides_parameter(
job_variables: Dict[str, Any], infra_overrides: Dict[str, Any]
Expand Down
15 changes: 15 additions & 0 deletions tests/_internal/compatibility/test_deprecated.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,3 +305,18 @@ def test_handles_setting_none(self, my_deployment_model):
json_dict = my_deployment.dict()
assert "job_variables" not in json_dict
assert json_dict["infra_overrides"] is None

def test_does_not_override_description_if_none_set(self, my_deployment_model):
assert "description" not in my_deployment_model.schema()

def test_preserves_description_if_present(self):
class MyModel(DeprecatedInfraOverridesField, pydantic.BaseModel):
"""Hey mom"""

name: str
job_variables: Optional[Dict[str, Any]] = pydantic.Field(
default_factory=dict,
description="Overrides to apply to my deployment infrastructure at runtime.",
)

assert MyModel.schema()["description"] == "Hey mom"

0 comments on commit ef09f72

Please sign in to comment.