chore: introduce helper mixins

Introduce FromDictMixin to provide a more type-safe way of deserializing a
dict into a class instance, and its reverse ToDictMixin.

Make model classes inherit the appropriate mixins.

Use the methods provided by the mixins instead of the not-type-safe
Marshmallow ones.

Author: agateau-gg

pygitguardian/client.py

@@ -10,12 +10,7 @@
 from requests import Response, Session, codes
 
 from .config import DEFAULT_API_VERSION, DEFAULT_BASE_URI, DEFAULT_TIMEOUT
-from .iac_models import (
-    IaCScanParameters,
-    IaCScanParametersSchema,
-    IaCScanResult,
-    IaCScanResultSchema,
-)
+from .iac_models import IaCScanParameters, IaCScanParametersSchema, IaCScanResult
 from .models import (
     Detail,
     Document,
@@ -65,7 +60,7 @@ def load_detail(resp: Response) -> Detail:
     else:
         data = {"detail": resp.text}
 
-    return cast(Detail, Detail.SCHEMA.load(data))
+    return Detail.from_dict(data)
 
 
 def is_ok(resp: Response) -> bool:
@@ -307,7 +302,7 @@ def content_scan(
         if filename:
             doc_dict["filename"] = filename
 
-        request_obj = Document.SCHEMA.load(doc_dict)
+        request_obj = cast(Dict[str, Any], Document.SCHEMA.load(doc_dict))
         DocumentSchema.validate_size(
             request_obj, self.secret_scan_preferences.maximum_document_size
         )
@@ -320,7 +315,7 @@ def content_scan(
 
         obj: Union[Detail, ScanResult]
         if is_ok(resp):
-            obj = ScanResult.SCHEMA.load(resp.json())
+            obj = ScanResult.from_dict(resp.json())
         else:
             obj = load_detail(resp)
 
@@ -354,7 +349,9 @@ def multi_content_scan(
             )
 
         if all(isinstance(doc, dict) for doc in documents):
-            request_obj = Document.SCHEMA.load(documents, many=True)
+            request_obj = cast(
+                List[Dict[str, Any]], Document.SCHEMA.load(documents, many=True)
+            )
         else:
             raise TypeError("each document must be a dict")
 
@@ -377,7 +374,7 @@ def multi_content_scan(
 
         obj: Union[Detail, MultiScanResult]
         if is_ok(resp):
-            obj = MultiScanResult.SCHEMA.load(dict(scan_results=resp.json()))
+            obj = MultiScanResult.from_dict({"scan_results": resp.json()})
         else:
             obj = load_detail(resp)
 
@@ -403,7 +400,7 @@ def quota_overview(
 
         obj: Union[Detail, QuotaResponse]
         if is_ok(resp):
-            obj = QuotaResponse.SCHEMA.load(resp.json())
+            obj = QuotaResponse.from_dict(resp.json())
         else:
             obj = load_detail(resp)
 
@@ -442,7 +439,7 @@ def create_honeytoken(
             result.status_code = 504
         else:
             if is_create_ok(resp):
-                result = HoneytokenResponse.SCHEMA.load(resp.json())
+                result = HoneytokenResponse.from_dict(resp.json())
             else:
                 result = load_detail(resp)
             result.status_code = resp.status_code
@@ -474,7 +471,7 @@ def iac_directory_scan(
             result.status_code = 504
         else:
             if is_ok(resp):
-                result = IaCScanResultSchema().load(resp.json())
+                result = IaCScanResult.from_dict(resp.json())
             else:
                 result = load_detail(resp)
 
@@ -497,7 +494,7 @@ def read_metadata(self) -> Optional[Detail]:
             result = load_detail(resp)
             result.status_code = resp.status_code
             return result
-        metadata = ServerMetadata.SCHEMA.load(resp.json())
+        metadata = ServerMetadata.from_dict(resp.json())
 
         self.secret_scan_preferences = metadata.secret_scan_preferences
         return None

pygitguardian/iac_models.py

@@ -1,13 +1,13 @@
 from dataclasses import dataclass, field
-from typing import List, Optional
+from typing import List, Optional, Type, cast
 
 import marshmallow_dataclass
 
-from pygitguardian.models import Base, BaseSchema
+from pygitguardian.models import Base, BaseSchema, FromDictMixin
 
 
 @dataclass
-class IaCVulnerability(Base):
+class IaCVulnerability(Base, FromDictMixin):
     policy: str
     policy_id: str
     line_end: int
@@ -18,37 +18,48 @@ class IaCVulnerability(Base):
     severity: str = ""
 
 
-IaCVulnerabilitySchema = marshmallow_dataclass.class_schema(
-    IaCVulnerability, BaseSchema
+IaCVulnerabilitySchema = cast(
+    Type[BaseSchema], marshmallow_dataclass.class_schema(IaCVulnerability, BaseSchema)
 )
 
+IaCVulnerability.SCHEMA = IaCVulnerabilitySchema()
+
 
 @dataclass
-class IaCFileResult(Base):
+class IaCFileResult(Base, FromDictMixin):
     filename: str
     incidents: List[IaCVulnerability]
 
 
-IaCFileResultSchema = marshmallow_dataclass.class_schema(IaCFileResult, BaseSchema)
+IaCFileResultSchema = cast(
+    Type[BaseSchema], marshmallow_dataclass.class_schema(IaCFileResult, BaseSchema)
+)
+
+IaCFileResult.SCHEMA = IaCFileResultSchema()
 
 
 @dataclass
-class IaCScanParameters(Base):
+class IaCScanParameters(Base, FromDictMixin):
     ignored_policies: List[str] = field(default_factory=list)
     minimum_severity: Optional[str] = None
 
 
-IaCScanParametersSchema = marshmallow_dataclass.class_schema(
-    IaCScanParameters, BaseSchema
+IaCScanParametersSchema = cast(
+    Type[BaseSchema], marshmallow_dataclass.class_schema(IaCScanParameters, BaseSchema)
 )
 
+IaCScanParameters.SCHEMA = IaCScanParametersSchema()
+
 
 @dataclass
-class IaCScanResult(Base):
+class IaCScanResult(Base, FromDictMixin):
     id: str = ""
     type: str = ""
     iac_engine_version: str = ""
     entities_with_incidents: List[IaCFileResult] = field(default_factory=list)
 
 
-IaCScanResultSchema = marshmallow_dataclass.class_schema(IaCScanResult, BaseSchema)
+IaCScanResultSchema = cast(
+    Type[BaseSchema], marshmallow_dataclass.class_schema(IaCScanResult, BaseSchema)
+)
+IaCScanResult.SCHEMA = IaCScanResultSchema()

pygitguardian/models.py

@@ -13,19 +13,48 @@
     pre_load,
     validate,
 )
+from typing_extensions import Self
 
 from .config import DOCUMENT_SIZE_THRESHOLD_BYTES, MULTI_DOCUMENT_LIMIT
 
 
+class ToDictMixin:
+    """
+    Provides a type-safe `to_dict()` method for classes using Marshmallow
+    """
+
+    SCHEMA: ClassVar[Schema]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return cast(Dict[str, Any], self.SCHEMA.dump(self))
+
+
+class FromDictMixin:
+    """This class must be used as an additional base class for all classes whose schema
+    implements a `post_load` function turning the received dict into a class instance.
+
+    It makes it possible to deserialize an object using `MyClass.from_dict(dct)` instead
+    of `MyClass.SCHEMA.load(dct)`. The `from_dict()` method is shorter, but more
+    importantly, type-safe: its return type is an instance of `MyClass`, not
+    `list[Any] | Any`.
+
+    Reference: https://marshmallow.readthedocs.io/en/stable/quickstart.html#deserializing-to-objects  E501
+    """
+
+    SCHEMA: ClassVar[Schema]
+
+    @classmethod
+    def from_dict(cls, dct: Dict[str, Any]) -> Self:
+        return cast(Self, cls.SCHEMA.load(dct))
+
+
 class BaseSchema(Schema):
     class Meta:
         ordered = True
         unknown = EXCLUDE
 
 
-class Base:
-    SCHEMA: ClassVar[BaseSchema]
-
+class Base(ToDictMixin):
     def __init__(self, status_code: Optional[int] = None) -> None:
         self.status_code = status_code
 
@@ -35,12 +64,6 @@ def to_json(self) -> str:
         """
         return cast(str, self.SCHEMA.dumps(self))
 
-    def to_dict(self) -> Dict:
-        """
-        to_dict converts model to a dictionary representation.
-        """
-        return cast(Dict, self.SCHEMA.dump(self))
-
     @property
     def success(self) -> bool:
         return self.__bool__()
@@ -122,7 +145,7 @@ def make_detail_response(self, data: Dict[str, Any], **kwargs: Any) -> "Detail":
         return Detail(**data)
 
 
-class Detail(Base):
+class Detail(Base, FromDictMixin):
     """Detail is a response object mostly returned on error or when the
     api output is a simple string.
 
@@ -155,7 +178,7 @@ def make_match(self, data: Dict[str, Any], **kwargs: Any) -> "Match":
         return Match(**data)
 
 
-class Match(Base):
+class Match(Base, FromDictMixin):
     """
     Match describes a found issue by GitGuardian.
     With info such as match location and type.
@@ -219,7 +242,7 @@ def make_policy_break(self, data: Dict[str, Any], **kwargs: Any) -> "PolicyBreak
         return PolicyBreak(**data)
 
 
-class PolicyBreak(Base):
+class PolicyBreak(Base, FromDictMixin):
     """
     PolicyBreak describes a GitGuardian policy break found
     in a scan.
@@ -269,7 +292,7 @@ def make_scan_result(self, data: Dict[str, Any], **kwargs: Any) -> "ScanResult":
         return ScanResult(**data)
 
 
-class ScanResult(Base):
+class ScanResult(Base, FromDictMixin):
     """ScanResult is a response object returned on a Content Scan
 
     Attributes:
@@ -355,7 +378,7 @@ def make_scan_result(
         return MultiScanResult(**data)
 
 
-class MultiScanResult(Base):
+class MultiScanResult(Base, FromDictMixin):
     """ScanResult is a response object returned on a Content Scan
 
     Attributes:
@@ -425,7 +448,7 @@ def make_quota(self, data: Dict[str, Any], **kwargs: Any) -> "Quota":
         return Quota(**data)
 
 
-class Quota(Base):
+class Quota(Base, FromDictMixin):
     """
     Quota describes a quota category in the GitGuardian API.
     Allows you to check your current available quota.
@@ -468,7 +491,7 @@ def make_quota_response(
         return QuotaResponse(**data)
 
 
-class QuotaResponse(Base):
+class QuotaResponse(Base, FromDictMixin):
     """
     Quota describes a quota category in the GitGuardian API.
     Allows you to check your current available quota.
@@ -515,7 +538,7 @@ def make_honeytoken_response(
         return HoneytokenResponse(**data)
 
 
-class HoneytokenResponse(Base):
+class HoneytokenResponse(Base, FromDictMixin):
     """
     honeytoken creation in the GitGuardian API.
     Allows users to create and get a honeytoken.
@@ -633,14 +656,15 @@ class SecretScanPreferences:
 
 
 @dataclass
-class ServerMetadata(Base):
+class ServerMetadata(Base, FromDictMixin):
     version: str
     preferences: Dict[str, Any]
     secret_scan_preferences: SecretScanPreferences = field(
         default_factory=SecretScanPreferences
     )
 
 
-ServerMetadata.SCHEMA = marshmallow_dataclass.class_schema(
-    ServerMetadata, base_schema=BaseSchema
-)()
+ServerMetadata.SCHEMA = cast(
+    BaseSchema,
+    marshmallow_dataclass.class_schema(ServerMetadata, base_schema=BaseSchema)(),
+)

setup.py

@@ -35,6 +35,7 @@ def get_version() -> str:
         "marshmallow>=3.5, <4",
         "requests>=2, <3",
         "marshmallow-dataclass >=8.5.8, <8.6.0",
+        "typing-extensions",
     ],
     include_package_data=True,
     zip_safe=True,