Refactor dump_json/dump_yaml into SerializableBaseModel base class

Author: NickCao

packages/jumpstarter-cli-common/jumpstarter_cli_common/version.py

@@ -3,10 +3,10 @@
 import sys
 
 import asyncclick as click
-import yaml
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import ConfigDict, Field
 
 from .opt import OutputMode, OutputType, opt_output_all
+from jumpstarter.common.pydantic import SerializableBaseModel
 
 
 def get_client_version():
@@ -27,18 +27,12 @@ def version_msg():
     return f"Jumpstarter v{jumpstarter_version} from {location} (Python {python_version})"
 
 
-class JumpstarterVersion(BaseModel):
+class JumpstarterVersion(SerializableBaseModel):
     git_version: str = Field(alias="gitVersion")
     python_version: str = Field(alias="pythonVersion")
 
     model_config = ConfigDict(populate_by_name=True)
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
-
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
-
 
 def version_obj():
     return JumpstarterVersion(git_version=importlib.metadata.version("jumpstarter"), python_version=sys.version)

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -3,13 +3,13 @@
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 
-import yaml
 from google.protobuf import duration_pb2, field_mask_pb2, json_format
 from grpc.aio import Channel
 from jumpstarter_protocol import client_pb2, client_pb2_grpc, kubernetes_pb2
-from pydantic import BaseModel, ConfigDict, Field, field_serializer
+from pydantic import ConfigDict, Field, field_serializer
 
 from jumpstarter.common.grpc import translate_grpc_exceptions
+from jumpstarter.common.pydantic import SerializableBaseModel
 
 
 def parse_identifier(identifier: str, kind: str) -> (str, str):
@@ -35,7 +35,7 @@ def parse_lease_identifier(identifier: str) -> (str, str):
     return parse_identifier(identifier, "leases")
 
 
-class Exporter(BaseModel):
+class Exporter(SerializableBaseModel):
     namespace: str
     name: str
     labels: dict[str, str]
@@ -46,7 +46,7 @@ def from_protobuf(cls, data: client_pb2.Exporter) -> Exporter:
         return cls(namespace=namespace, name=name, labels=data.labels)
 
 
-class Lease(BaseModel):
+class Lease(SerializableBaseModel):
     namespace: str
     name: str
     selector: str
@@ -92,14 +92,8 @@ def from_protobuf(cls, data: client_pb2.Lease) -> Lease:
             conditions=data.conditions,
         )
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
 
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
-
-
-class ExporterList(BaseModel):
+class ExporterList(SerializableBaseModel):
     exporters: list[Exporter]
     next_page_token: str | None = Field(exclude=True)
 
@@ -110,14 +104,8 @@ def from_protobuf(cls, data: client_pb2.ListExportersResponse) -> ExporterList:
             next_page_token=data.next_page_token,
         )
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
-
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
 
-
-class LeaseList(BaseModel):
+class LeaseList(SerializableBaseModel):
     leases: list[Lease]
     next_page_token: str | None = Field(exclude=True)
 
@@ -128,12 +116,6 @@ def from_protobuf(cls, data: client_pb2.ListLeasesResponse) -> LeaseList:
             next_page_token=data.next_page_token,
         )
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
-
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
-
 
 @dataclass(kw_only=True, slots=True)
 class ClientService:

packages/jumpstarter/jumpstarter/common/pydantic.py

@@ -0,0 +1,10 @@
+import yaml
+from pydantic import BaseModel
+
+
+class SerializableBaseModel(BaseModel):
+    def dump_json(self):
+        return self.model_dump_json(indent=4, by_alias=True)
+
+    def dump_yaml(self):
+        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)

packages/jumpstarter/jumpstarter/config/client.py

@@ -16,6 +16,7 @@
 from jumpstarter.client.grpc import ClientService
 from jumpstarter.common.exceptions import FileNotFoundError
 from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
+from jumpstarter.common.pydantic import SerializableBaseModel
 
 
 def _allow_from_env():
@@ -274,16 +275,10 @@ def delete(cls, alias: str) -> Path:
         return path
 
 
-class ClientConfigListV1Alpha1(BaseModel):
+class ClientConfigListV1Alpha1(SerializableBaseModel):
     api_version: Literal["jumpstarter.dev/v1alpha1"] = Field(alias="apiVersion", default="jumpstarter.dev/v1alpha1")
     current_config: Optional[str] = Field(alias="currentConfig")
     items: list[ClientConfigV1Alpha1]
     kind: Literal["ClientConfigList"] = Field(default="ClientConfigList")
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
-
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
-
     model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)

packages/jumpstarter/jumpstarter/config/exporter.py

@@ -14,6 +14,7 @@
 from .tls import TLSConfigV1Alpha1
 from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
 from jumpstarter.common.importlib import import_class
+from jumpstarter.common.pydantic import SerializableBaseModel
 from jumpstarter.driver import Driver
 
 
@@ -175,15 +176,9 @@ async def channel_factory():
             await exporter.serve()
 
 
-class ExporterConfigListV1Alpha1(BaseModel):
+class ExporterConfigListV1Alpha1(SerializableBaseModel):
     api_version: Literal["jumpstarter.dev/v1alpha1"] = Field(alias="apiVersion", default="jumpstarter.dev/v1alpha1")
     items: list[ExporterConfigV1Alpha1]
     kind: Literal["ExporterConfigList"] = Field(default="ExporterConfigList")
 
-    def dump_json(self):
-        return self.model_dump_json(indent=4, by_alias=True)
-
-    def dump_yaml(self):
-        return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
-
     model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)