Skip to content

Commit

Permalink
fix: pinned bool columns show up as None (#225)
Browse files Browse the repository at this point in the history
* dont generalize falsiness

* add test
  • Loading branch information
isabelizimm authored Dec 10, 2024
1 parent f3e53e6 commit a466e36
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,41 @@ def test_vetiver_model_dict_like_prototype(prototype_data):
assert json_schema == expected


@pytest.mark.skipif(
pydantic.__version__.startswith("1"), reason="only run for pydantic v2"
)
@pytest.mark.parametrize(
"prototype_data,expected",
[
(
{"B": 0, "C": False, "D": None},
{
"properties": {
"B": {"example": 0, "title": "B", "type": "integer"},
"C": {"example": False, "title": "C", "type": "boolean"},
"D": {"example": None, "title": "D", "type": "null"},
},
"required": ["B", "C", "D"],
"title": "prototype",
"type": "object",
},
)
],
)
def test_falsy_prototypes(prototype_data, expected):
v = VetiverModel(
model=model,
prototype_data=prototype_data,
model_name="model",
versioned=None,
description=None,
metadata=None,
)

assert isinstance(v.prototype.construct(), pydantic.BaseModel)
assert v.prototype.model_json_schema() == expected


@pytest.mark.parametrize("prototype_data", [MockPrototype(B=4, C=0, D=0), None])
def test_vetiver_model_prototypes(prototype_data):
v = VetiverModel(
Expand Down
4 changes: 3 additions & 1 deletion vetiver/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,8 @@ def serialize_prototype(prototype):

serialized_schema = dict()
for key, value in schema.items():
serialized_schema[key] = value.get("example") or value.get("default")
example = value.get("example", None)
default = value.get("default", None)
serialized_schema[key] = example if example is not None else default

return json.dumps(serialized_schema)

0 comments on commit a466e36

Please sign in to comment.