Instate msgspec type registration mechanism (#199)

* Instate msgspec type registration mechanism

* Update ChangeSet API

* address review

Author: ofek

docs/reference/api/git.md

@@ -34,13 +34,12 @@
 ::: dda.utils.git.changeset.ChangeSet
     options:
       members:
-      - changes
+      - files
+      - paths
       - added
       - modified
       - deleted
-      - changed
       - digest
-      - from_iter
       - from_patches
 
 ::: dda.utils.git.changeset.ChangedFile

src/dda/config/model/__init__.py

@@ -15,7 +15,7 @@
 from dda.config.model.tools import ToolsConfig
 from dda.config.model.update import UpdateConfig
 from dda.config.model.user import UserConfig
-from dda.utils.fs import Path
+from dda.types.hooks import dec_hook, enc_hook
 
 
 def _default_orgs() -> dict[str, OrgConfig]:
@@ -52,38 +52,3 @@ def get_default_toml_data() -> dict[str, Any]:
         builtin_types=(datetime.datetime, datetime.date, datetime.time),
         enc_hook=enc_hook,
     )
-
-
-def dec_hook(type: type[Any], obj: Any) -> Any:  # noqa: A002
-    if type is Path:
-        return Path(obj)
-
-    from msgspec import convert
-
-    from dda.utils.git.changeset import ChangedFile, ChangeSet
-
-    if type is ChangeSet:
-        # Since the dict decode logic from msgspec is not called here we have to manually decode the keys and values
-        decoded_obj = {}
-        for key, value in obj.items():
-            decoded_key = dec_hook(Path, key)
-            decoded_value = convert(value, ChangedFile, dec_hook=dec_hook)
-            decoded_obj[decoded_key] = decoded_value
-        return ChangeSet(changes=decoded_obj)
-
-    message = f"Cannot decode: {obj!r}"
-    raise ValueError(message)
-
-
-def enc_hook(obj: Any) -> Any:
-    if isinstance(obj, Path):
-        return str(obj)
-
-    from dda.utils.git.changeset import ChangeSet
-
-    # Encode ChangeSet objects as dicts
-    if isinstance(obj, ChangeSet):
-        return dict(obj.changes)
-
-    message = f"Cannot encode: {obj!r}"
-    raise NotImplementedError(message)

src/dda/types/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2025-present Datadog, Inc. <dev@datadoghq.com>
+#
+# SPDX-License-Identifier: MIT

src/dda/types/hooks.py

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: 2025-present Datadog, Inc. <dev@datadoghq.com>
+#
+# SPDX-License-Identifier: MIT
+from __future__ import annotations
+
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any
+
+from msgspec import Struct
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class Hook(Struct, frozen=True):
+    encode: Callable[[Any], Any]
+    decode: Callable[[Any], Any]
+
+
+def register_type_hooks(
+    typ: type[Any],
+    *,
+    encode: Callable[[Any], Any],
+    decode: Callable[[Any], Any],
+) -> None:
+    __HOOKS[typ] = Hook(encode=encode, decode=decode)
+
+
+def enc_hook(obj: Any) -> Any:
+    if (registered_type := __HOOKS.get(type(obj))) is not None:
+        return registered_type.encode(obj)
+
+    message = f"Cannot encode: {obj!r}"
+    raise NotImplementedError(message)
+
+
+def dec_hook(typ: type[Any], obj: Any) -> Any:
+    if (registered_type := __HOOKS.get(typ)) is not None:
+        return registered_type.decode(obj)
+
+    message = f"Cannot decode: {obj!r}"
+    raise ValueError(message)
+
+
+__HOOKS: dict[type[Any], Hook] = {}
+
+register_type_hooks(MappingProxyType, encode=dict, decode=MappingProxyType)

src/dda/utils/fs.py

@@ -10,6 +10,8 @@
 from functools import cached_property
 from typing import IO, TYPE_CHECKING, Any
 
+from dda.types.hooks import register_type_hooks
+
 if TYPE_CHECKING:
     from collections.abc import Generator
 
@@ -189,3 +191,6 @@ def temp_file(suffix: str = "") -> Generator[Path, None, None]:
 
     with NamedTemporaryFile(suffix=suffix) as f:
         yield Path(f.name).resolve()
+
+
+register_type_hooks(Path, encode=str, decode=Path)

src/dda/utils/git/changeset.py

@@ -6,15 +6,17 @@
 
 from enum import StrEnum
 from functools import cached_property
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Self
 
-from msgspec import Struct
+from msgspec import Struct, convert, to_builtins
 
+from dda.types.hooks import dec_hook, enc_hook, register_type_hooks
 from dda.utils.fs import Path
 
 if TYPE_CHECKING:
-    from collections.abc import ItemsView, Iterable, Iterator, KeysView, ValuesView
+    from collections.abc import Iterable
 
 
 class ChangeType(StrEnum):
@@ -58,78 +60,45 @@ class ChangeSet:  # noqa: PLW1641
     When considering the changes to the working directory, the untracked files are considered as added files.
     """
 
-    def __init__(self, changes: dict[Path, ChangedFile]) -> None:
-        self.__changes = MappingProxyType(changes)
+    def __init__(self, changed_files: Iterable[ChangedFile]) -> None:
+        self.__changed = MappingProxyType({str(c.path): c for c in changed_files})
+        self.__files = tuple(self.__changed.values())
 
     @property
-    def changes(self) -> MappingProxyType[Path, ChangedFile]:
-        return self.__changes
+    def paths(self) -> MappingProxyType[str, ChangedFile]:
+        return self.__changed
 
-    def keys(self) -> KeysView[Path]:
-        return self.__changes.keys()
-
-    def values(self) -> ValuesView[ChangedFile]:
-        return self.__changes.values()
-
-    def items(self) -> ItemsView[Path, ChangedFile]:
-        return self.__changes.items()
-
-    def __getitem__(self, key: Path) -> ChangedFile:
-        return self.__changes[key]
-
-    def __contains__(self, key: Path) -> bool:
-        return key in self.__changes
-
-    def __len__(self) -> int:
-        return len(self.__changes)
-
-    def __iter__(self) -> Iterator[Path]:
-        return iter(self.__changes.keys())
-
-    def __or__(self, other: Self) -> Self:
-        return self.from_iter(list(self.values()) + list(other.values()))
-
-    def __eq__(self, other: object) -> bool:
-        return isinstance(other, ChangeSet) and self.__changes == other.__changes
-
-    @cached_property
-    def added(self) -> set[Path]:
-        """List of files that were added."""
-        return {change.path for change in self.values() if change.type == ChangeType.ADDED}
+    @property
+    def files(self) -> Iterable[ChangedFile]:
+        return self.__files
 
-    @cached_property
-    def modified(self) -> set[Path]:
-        """List of files that were modified."""
-        return {change.path for change in self.values() if change.type == ChangeType.MODIFIED}
+    @property
+    def added(self) -> MappingProxyType[str, ChangedFile]:
+        """Set of files that were added."""
+        return self.__change_types[ChangeType.ADDED]
 
-    @cached_property
-    def deleted(self) -> set[Path]:
-        """List of files that were deleted."""
-        return {change.path for change in self.values() if change.type == ChangeType.DELETED}
+    @property
+    def modified(self) -> MappingProxyType[str, ChangedFile]:
+        """Set of files that were modified."""
+        return self.__change_types[ChangeType.MODIFIED]
 
-    @cached_property
-    def changed(self) -> set[Path]:
-        """List of files that were changed (added, modified, or deleted)."""
-        return set(self.keys())
+    @property
+    def deleted(self) -> MappingProxyType[str, ChangedFile]:
+        """Set of files that were deleted."""
+        return self.__change_types[ChangeType.DELETED]
 
     def digest(self) -> str:
         """Compute a hash of the changeset."""
         from hashlib import sha256
 
         digester = sha256()
-        for change in sorted(self.values(), key=lambda x: x.path.as_posix()):
+        for change in sorted(self.files, key=lambda cf: cf.path):
             digester.update(change.path.as_posix().encode())
             digester.update(change.type.value.encode())
             digester.update(change.patch.encode())
 
         return str(digester.hexdigest())
 
-    @classmethod
-    def from_iter(cls, data: Iterable[ChangedFile]) -> Self:
-        """Create a ChangeSet from an iterable of FileChanges."""
-        items = {change.path: change for change in data}
-        return cls(changes=items)
-
     @classmethod
     def from_patches(cls, diff_output: str | list[str]) -> Self:
         """
@@ -182,7 +151,21 @@ def from_patches(cls, diff_output: str | list[str]) -> Self:
             # Strip every "block" and add the missing separator
             patch = "" if binary else "\n".join([sep + block.strip() for block in blocks]).strip()
             changes.append(ChangedFile(path=current_file, type=current_type, binary=binary, patch=patch))
-        return cls.from_iter(changes)
+        return cls(changes)
+
+    def __or__(self, other: Self) -> Self:
+        return type(self)(chain(self.files, other.files))
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, ChangeSet) and self.paths == other.paths
+
+    @cached_property
+    def __change_types(self) -> dict[ChangeType, MappingProxyType[str, ChangedFile]]:
+        changes: dict[ChangeType, dict[str, ChangedFile]] = {}
+        for change in self.files:
+            changes.setdefault(change.type, {})[str(change.path)] = change
+
+        return {change_type: MappingProxyType(paths) for change_type, paths in changes.items()}
 
 
 def _determine_change_type(before_filename: str, after_filename: str) -> ChangeType:
@@ -195,3 +178,10 @@ def _determine_change_type(before_filename: str, after_filename: str) -> ChangeT
 
     msg = f"Unexpected file paths in git diff output: {before_filename} -> {after_filename} - this indicates a rename which we do not support"
     raise ValueError(msg)
+
+
+register_type_hooks(
+    ChangeSet,
+    encode=lambda obj: to_builtins(obj.files, enc_hook=enc_hook),
+    decode=lambda obj: ChangeSet(convert(cf, ChangedFile, dec_hook=dec_hook) for cf in obj),
+)

src/dda/utils/git/github.py

@@ -44,7 +44,7 @@ def get_commit_and_changes_from_github(remote: Remote, sha1: str) -> tuple[Commi
     data = client.get(get_commit_github_api_url(remote, sha1)).json()
 
     # Compute ChangeSet
-    changes = ChangeSet.from_iter(
+    changes = ChangeSet(
         ChangedFile(
             path=Path(file_obj["filename"]),
             type=get_change_type_from_github_status(file_obj["status"]),

tests/tools/git/fixtures/repo_states/complex_changes/expected_changeset.json

@@ -1,32 +1,32 @@
-{
-    "file2.txt": {
+[
+    {
         "path": "file2.txt",
         "type": "D",
         "binary": false,
         "patch": "@@ -1,4 +0,0 @@\n-I am file2 !\n-I feel like I take up space for nothing...\n-I have a feeling like I won't exist pretty soon :/\n-"
     },
-    "file3.txt": {
+    {
         "path": "file3.txt",
         "type": "A",
         "binary": false,
         "patch": "@@ -0,0 +1,2 @@\n+I am file3 !\n+I'm new around here, hopefully everyone treats me nice :)"
     },
-    "file4.txt": {
+    {
         "path": "file4.txt",
         "type": "M",
         "binary": false,
         "patch": "@@ -2 +2 @@ I am file4.\n-People often tell me I am unreliable.\n+People often tell me I am THE BEST.\n@@ -4,3 +4,2 @@ Things like:\n-- You always change !\n-- I can never count on you...\n-- I didn't recognize you !\n+- You rock !\n+- I wish I were you !\n@@ -8 +7,3 @@ Do you think they have a point ?\n-I'd need to look at my own history to know...\n+Pah ! Who am I kidding, they're OBVIOUSLY RIGHT.\n+Arrogance ? What is that, an italian ice cream flavor ?\n+Get outta here !"
     },
-    "file5.txt": {
+    {
         "path": "file5.txt",
         "type": "D",
         "binary": false,
         "patch": "@@ -1,5 +0,0 @@\n-I am a humble file.\n-Soon I will change name.\n-I think I'll also take this as an opportunity to change myself.\n-New name, new me !\n-Or is that not how the saying goes ?"
     },
-    "file5_new.txt": {
+    {
         "path": "file5_new.txt",
         "type": "A",
         "binary": false,
         "patch": "@@ -0,0 +1,5 @@\n+I am a humble file.\n+Hey I have a new name !\n+Wow, I feel much better now.\n+New name, new me !\n+Or is that not how the saying goes ?"
     }
-}
+]

tests/tools/git/fixtures/repo_states/file_add/expected_changeset.json

@@ -1,8 +1,8 @@
-{
-    "file2.txt": {
+[
+    {
         "path": "file2.txt",
         "type": "A",
         "binary": false,
         "patch": "@@ -0,0 +1,3 @@\n+file2\n+I am a new file in the repo !\n+That's incredible."
     }
-}
+]

tests/tools/git/fixtures/repo_states/file_delete/expected_changeset.json

@@ -1,8 +1,8 @@
-{
-    "file2.txt": {
+[
+    {
         "path": "file2.txt",
         "type": "D",
         "binary": false,
         "patch": "@@ -1,3 +0,0 @@\n-file2\n-I will be deleted, unfortunately.\n-That's quite sad."
     }
-}
+]