Merge pull request #155 from SinaKhalili/sina/add-types-coder

Add types coder

Author: kevinheavey

src/anchorpy/coder/coder.py

@@ -4,6 +4,7 @@
 from anchorpy.coder.accounts import AccountsCoder
 from anchorpy.coder.event import EventCoder
 from anchorpy.coder.instruction import InstructionCoder
+from anchorpy.coder.types import TypesCoder
 
 
 class Coder:
@@ -18,3 +19,4 @@ def __init__(self, idl: Idl):
         self.instruction: InstructionCoder = InstructionCoder(idl)
         self.accounts: AccountsCoder = AccountsCoder(idl)
         self.events: EventCoder = EventCoder(idl)
+        self.types: TypesCoder = TypesCoder(idl)

src/anchorpy/coder/types.py

@@ -0,0 +1,82 @@
+"""The TypesCoder class is for encoding and decoding user-defined types."""
+
+from typing import Any, Dict
+
+from anchorpy_core.idl import Idl
+from construct import Construct, Container
+
+from anchorpy.coder.idl import _typedef_layout_without_field_name
+
+
+class TypesCoder:
+    """Encodes and decodes user-defined types in Anchor programs."""
+
+    def __init__(self, idl: Idl) -> None:
+        """Initialize the TypesCoder.
+
+        Args:
+            idl: The parsed IDL object.
+        """
+        self.idl = idl
+        self.types_layouts: Dict[str, Construct] = {}
+
+        self.filtered_types = []
+        if idl.types:
+            self.filtered_types = [
+                ty for ty in idl.types if not getattr(ty, "generics", None)
+            ]
+
+    def _get_layout(self, name: str) -> Construct:
+        """Get or create a layout for a given type name.
+
+        Args:
+            name: The name of the type.
+
+        Returns:
+            The construct layout for the type.
+
+        Raises:
+            ValueError: If the type is not found.
+        """
+        if name in self.types_layouts:
+            return self.types_layouts[name]
+
+        type_defs = [t for t in self.filtered_types if t.name == name]
+        if not type_defs:
+            raise ValueError(f"Unknown type: {name}")
+
+        layout = _typedef_layout_without_field_name(type_defs[0], self.idl.types)
+        self.types_layouts[name] = layout
+        return layout
+
+    def encode(self, name: str, data: Any) -> bytes:
+        """Encode a user-defined type.
+
+        Args:
+            name: The name of the type.
+            data: The data to encode.
+
+        Returns:
+            The encoded data.
+
+        Raises:
+            ValueError: If the type is not found.
+        """
+        layout = self._get_layout(name)
+        return layout.build(data)
+
+    def decode(self, name: str, buffer: bytes) -> Container[Any]:
+        """Decode a user-defined type.
+
+        Args:
+            name: The name of the type.
+            buffer: The buffer to decode.
+
+        Returns:
+            The decoded data.
+
+        Raises:
+            ValueError: If the type is not found.
+        """
+        layout = self._get_layout(name)
+        return layout.parse(buffer)

tests/unit/test_types_coder.py

@@ -0,0 +1,102 @@
+import json
+
+import pytest
+from anchorpy import Coder
+from anchorpy_core.idl import Idl
+
+
+@pytest.mark.unit
+def test_can_encode_and_decode_user_defined_types():
+    """Test that the TypesCoder can encode and decode user-defined types."""
+    idl_json = {
+        "version": "0.0.0",
+        "name": "basic_0",
+        "address": "Test111111111111111111111111111111111111111",
+        "instructions": [
+            {
+                "name": "initialize",
+                "accounts": [],
+                "args": [],
+                "discriminator": [],
+            },
+        ],
+        "types": [
+            {
+                "name": "MintInfo",
+                "type": {
+                    "kind": "struct",
+                    "fields": [
+                        {
+                            "name": "minted",
+                            "type": "bool",
+                        },
+                        {
+                            "name": "metadataUrl",
+                            "type": "string",
+                        },
+                    ],
+                },
+            },
+        ],
+    }
+    idl = Idl.from_json(json.dumps(idl_json))
+    coder = Coder(idl)
+
+    mint_info = {
+        "minted": True,
+        "metadata_url": "hello",
+    }
+    encoded = coder.types.encode("MintInfo", mint_info)
+    decoded = coder.types.decode("MintInfo", encoded)
+
+    # Compare decoded values with original
+    assert decoded.minted == mint_info["minted"]
+    assert decoded.metadata_url == mint_info["metadata_url"]
+
+
+@pytest.mark.unit
+def test_can_encode_and_decode_large_integers():
+    """Test that the TypesCoder can encode and decode 128-bit integers."""
+    idl_json = {
+        "version": "0.0.0",
+        "name": "basic_0",
+        "address": "Test111111111111111111111111111111111111111",
+        "instructions": [
+            {
+                "name": "initialize",
+                "accounts": [],
+                "args": [],
+                "discriminator": [],
+            },
+        ],
+        "types": [
+            {
+                "name": "IntegerTest",
+                "type": {
+                    "kind": "struct",
+                    "fields": [
+                        {
+                            "name": "unsigned",
+                            "type": "u128",
+                        },
+                        {
+                            "name": "signed",
+                            "type": "i128",
+                        },
+                    ],
+                },
+            },
+        ],
+    }
+    idl = Idl.from_json(json.dumps(idl_json))
+    coder = Coder(idl)
+
+    integer_test = {
+        "unsigned": 2588012355,
+        "signed": -93842345,
+    }
+    encoded = coder.types.encode("IntegerTest", integer_test)
+    decoded = coder.types.decode("IntegerTest", encoded)
+
+    assert decoded.unsigned == integer_test["unsigned"]
+    assert decoded.signed == integer_test["signed"]