Skip to content

Commit

Permalink
fix: safe dump bento info
Browse files Browse the repository at this point in the history
Signed-off-by: Frost Ming <me@frostming.com>
  • Loading branch information
frostming committed Jan 26, 2024
1 parent e524a9d commit a658e64
Showing 1 changed file with 2 additions and 14 deletions.
16 changes: 2 additions & 14 deletions src/bentoml/_internal/bento/bento.py
Expand Up @@ -626,9 +626,9 @@ def __attrs_post_init__(self):
def to_dict(self) -> t.Dict[str, t.Any]:
return bentoml_cattr.unstructure(self)

def dump(self, stream: t.IO[t.Any]):
def dump(self, stream: t.IO[t.Any]) -> None:
# _models is an alias for models, replace it with models
return yaml.dump(self, stream, sort_keys=False)
return yaml.safe_dump(self.to_dict(), stream, sort_keys=False)

@classmethod
def from_yaml_file(cls, stream: t.IO[t.Any]) -> BentoInfo:
Expand Down Expand Up @@ -698,15 +698,3 @@ def validate(self):
# Ignore tag, tag is saved via the name and version field
make_dict_unstructure_fn(BentoInfo, bentoml_cattr, tag=override(omit=True)),
)


def _BentoInfo_dumper(dumper: yaml.Dumper, info: BentoInfo) -> yaml.Node:
return dumper.represent_dict(info.to_dict())


def _tuple_dumper(dumper: yaml.Dumper, info: tuple[t.Any, ...]) -> yaml.Node:
return dumper.represent_list(info)


yaml.add_representer(BentoInfo, _BentoInfo_dumper) # type: ignore (incomplete yaml types)
yaml.add_representer(tuple, _tuple_dumper) # type: ignore (incomplete yaml types)

0 comments on commit a658e64

Please sign in to comment.