Add more pydantic utilities (#148)

Author: cthoyt

src/pystow/utils/__init__.py

@@ -82,12 +82,15 @@
     ModelValidateFailureAction,
     iter_pydantic_jsonl,
     iter_pydantic_tsv,
+    model_dump_yaml,
     read_pydantic_json,
     read_pydantic_jsonl,
     read_pydantic_tsv,
     read_pydantic_yaml,
     stream_write_pydantic_jsonl,
+    write_pydantic_json,
     write_pydantic_jsonl,
+    write_pydantic_yaml,
 )
 from .safe_open import (
     is_url,
@@ -97,6 +100,8 @@
     safe_open_dict_reader,
     safe_open_json,
     safe_open_yaml,
+    write_json,
+    write_yaml,
 )
 from ..constants import README_TEXT, TimeoutHint
 
@@ -153,6 +158,7 @@
     "mkdir",
     "mock_envvar",
     "mock_home",
+    "model_dump_yaml",
     "n",
     "name_from_s3_key",
     "name_from_url",
@@ -188,11 +194,15 @@
     "stream_write_pydantic_jsonl",
     "tarfile_writestr",
     "use_appdirs",
+    "write_json",
     "write_lzma_csv",
     "write_pickle_gz",
+    "write_pydantic_json",
     "write_pydantic_jsonl",
+    "write_pydantic_yaml",
     "write_tarfile_csv",
     "write_tarfile_xml",
+    "write_yaml",
     "write_zipfile_csv",
     "write_zipfile_np",
     "write_zipfile_rdf",

src/pystow/utils/pydantic_utils.py

@@ -10,7 +10,15 @@
 
 from tqdm import tqdm
 
-from .safe_open import safe_open, safe_open_dict_reader, safe_open_json, safe_open_yaml
+from .safe_open import (
+    _open_read_text,
+    _open_write_text,
+    safe_open_dict_reader,
+    safe_open_json,
+    safe_open_yaml,
+    write_json,
+    write_yaml,
+)
 
 if TYPE_CHECKING:
     import pydantic
@@ -19,12 +27,15 @@
     "ModelValidateFailureAction",
     "iter_pydantic_jsonl",
     "iter_pydantic_tsv",
+    "model_dump_yaml",
     "read_pydantic_json",
     "read_pydantic_jsonl",
     "read_pydantic_tsv",
     "read_pydantic_yaml",
     "stream_write_pydantic_jsonl",
+    "write_pydantic_json",
     "write_pydantic_jsonl",
+    "write_pydantic_yaml",
 ]
 
 logger = logging.getLogger(__name__)
@@ -55,9 +66,7 @@ def iter_pydantic_jsonl(
     }
     if tqdm_kwargs is not None:
         _tqdm_kwargs.update(tqdm_kwargs)
-    with safe_open(
-        file, operation="read", representation="text", encoding=encoding, newline=newline
-    ) as file:
+    with _open_read_text(file, encoding=encoding, newline=newline) as file:
         for i, line in enumerate(tqdm(file, disable=not progress, **_tqdm_kwargs)):
             try:
                 yv = model_cls.model_validate_json(line.strip())
@@ -85,7 +94,7 @@ def write_pydantic_jsonl(
     kwargs.setdefault("exclude_none", True)
     kwargs.setdefault("exclude_unset", True)
     kwargs.setdefault("exclude_defaults", True)
-    with safe_open(file, operation="write", representation="text") as file:
+    with _open_write_text(file) as file:
         for model in models:
             file.write(model.model_dump_json(**kwargs) + "\n")
 
@@ -97,7 +106,7 @@ def stream_write_pydantic_jsonl(
     kwargs.setdefault("exclude_none", True)
     kwargs.setdefault("exclude_unset", True)
     kwargs.setdefault("exclude_defaults", True)
-    with safe_open(file, operation="write", representation="text") as file:
+    with _open_write_text(file) as file:
         for model in models:
             file.write(model.model_dump_json(**kwargs) + "\n")
             yield model
@@ -161,3 +170,44 @@ def read_pydantic_yaml(
 ) -> BaseModelVar:
     """Read a YAML file into a model."""
     return model_cls.model_validate(safe_open_yaml(path_or_url, encoding=encoding, newline=newline))
+
+
+def model_dump_yaml(
+    model: pydantic.BaseModel,
+    *,
+    exclude_none: bool = False,
+    exclude_unset: bool = False,
+) -> str:
+    """Dump the model as YAML string."""
+    import yaml
+
+    data = model.model_dump(mode="json", exclude_none=exclude_none, exclude_unset=exclude_unset)
+    return yaml.safe_dump(data)
+
+
+def write_pydantic_yaml(
+    model: pydantic.BaseModel,
+    path: str | Path | TextIO,
+    *,
+    exclude_none: bool = False,
+    exclude_unset: bool = False,
+    encoding: str | None = None,
+    newline: str | None = None,
+) -> None:
+    """Write a model to a YAML file."""
+    data = model_dump_yaml(model, exclude_none=exclude_none, exclude_unset=exclude_unset)
+    write_yaml(data, path, encoding=encoding, newline=newline)
+
+
+def write_pydantic_json(
+    model: pydantic.BaseModel,
+    path: str | Path | TextIO,
+    *,
+    exclude_none: bool = False,
+    exclude_unset: bool = False,
+    encoding: str | None = None,
+    newline: str | None = None,
+) -> None:
+    """Write a model to a JSON file."""
+    data = model.model_dump(mode="json", exclude_none=exclude_none, exclude_unset=exclude_unset)
+    write_json(data, path, encoding=encoding, newline=newline)

src/pystow/utils/safe_open.py

@@ -35,6 +35,8 @@
     "safe_open_dict_reader",
     "safe_open_json",
     "safe_open_yaml",
+    "write_json",
+    "write_yaml",
 ]
 
 
@@ -142,16 +144,38 @@ def safe_open(  # noqa:C901
         raise TypeError(f"unsupported type for opening: {type(path)} - {path}")
 
 
+@contextlib.contextmanager
+def _open_read_text(
+    path: str | Path | TextIO,
+    encoding: str | None = None,
+    newline: str | None = None,
+) -> Generator[typing.TextIO, None, None]:
+    with safe_open(
+        path, representation="text", operation="read", encoding=encoding, newline=newline
+    ) as file:
+        yield file
+
+
+@contextlib.contextmanager
+def _open_write_text(
+    path: str | Path | TextIO,
+    encoding: str | None = None,
+    newline: str | None = None,
+) -> Generator[typing.TextIO, None, None]:
+    with safe_open(
+        path, representation="text", operation="write", encoding=encoding, newline=newline
+    ) as file:
+        yield file
+
+
 def safe_open_json(
     path_or_url: str | Path | TextIO,
     *,
     encoding: str | None = None,
     newline: str | None = None,
 ) -> Any:
     """Safely open a file and parse as JSON."""
-    with safe_open(
-        path_or_url, representation="text", operation="read", encoding=encoding, newline=newline
-    ) as file:
+    with _open_read_text(path_or_url, encoding=encoding, newline=newline) as file:
         return json.load(file)
 
 
@@ -164,12 +188,38 @@ def safe_open_yaml(
     """Safely open a file and parse as YAML."""
     import yaml
 
-    with safe_open(
-        path_or_url, representation="text", operation="read", encoding=encoding, newline=newline
-    ) as file:
+    with _open_read_text(path_or_url, encoding=encoding, newline=newline) as file:
         return yaml.safe_load(file)
 
 
+def write_yaml(
+    data: Any,
+    path: str | Path | TextIO,
+    *,
+    encoding: str | None = None,
+    newline: str | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Write YAML to a file."""
+    import yaml
+
+    with _open_write_text(path, encoding=encoding, newline=newline) as file:
+        yaml.safe_dump(data, file, **kwargs)
+
+
+def write_json(
+    data: Any,
+    path: str | Path | TextIO,
+    *,
+    encoding: str | None = None,
+    newline: str | None = None,
+    **kwargs: Any,
+) -> Any:
+    """Write JSON to a file."""
+    with _open_write_text(path, encoding=encoding, newline=newline) as file:
+        json.dump(data, file, **kwargs)
+
+
 # docstr-coverage:excused `overload`
 @typing.overload
 @contextlib.contextmanager
@@ -237,7 +287,7 @@ def safe_open_dict_reader(
 
     :yields: A CSV reader object, constructed from :func:`csv.DictReader`
     """
-    with safe_open(f, operation="read", representation="text") as file:
+    with _open_read_text(f) as file:
         yield csv.DictReader(file, delimiter=delimiter, **kwargs)