Fix parsing of BlobRef in CBOR

MarshalX · Dec 21, 2023 · 8948077 · 8948077
1 parent 6de3060
commit 8948077
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 9 deletions.
diff --git a/packages/atproto_client/models/blob_ref.py b/packages/atproto_client/models/blob_ref.py
@@ -1,16 +1,14 @@
 import typing as t
 
 import typing_extensions as te
+from atproto_core.cid import CID
 from pydantic import BaseModel, ConfigDict, Field
 
-if t.TYPE_CHECKING:
-    from atproto_core.cid import CID
 
+class IpldLink(BaseModel):
+    """CID representation in JSON."""
 
-class BlobRefLink(BaseModel):
-    """Blob reference link."""
-
-    link: str = Field(alias='$link')
+    link: str = Field(alias='$link')  #: CID.
 
 
 class BlobRef(BaseModel):
@@ -20,13 +18,14 @@ class BlobRef(BaseModel):
 
     mime_type: str = Field(alias='mimeType')  #: Mime type.
     size: int  #: Size in bytes.
-    ref: BlobRefLink  #: Reference to link.
+    ref: t.Union[str, IpldLink]  #: CID.
 
     py_type: te.Literal['blob'] = Field(default='blob', alias='$type')
 
     @property
     def cid(self) -> 'CID':
         """Get CID."""
-        from atproto_core.cid import CID
+        if isinstance(self.ref, str):
+            return CID.decode(self.ref)
 
         return CID.decode(self.ref.link)
diff --git a/packages/atproto_core/cid/cid.py b/packages/atproto_core/cid/cid.py
@@ -15,7 +15,7 @@ class Multihash:
     digest: bytes
 
 
-@dataclass
+@dataclass(eq=False)
 class CID:
     _cid: str
     version: int
@@ -48,6 +48,15 @@ def __str__(self) -> str:
     def __hash__(self) -> int:
         return hash(self.encode())
 
+    def __eq__(self, other: t.Any) -> bool:
+        if isinstance(other, str):
+            return self.encode() == other
+
+        if isinstance(other, CID):
+            return self.encode() == other.encode()
+
+        return False
+
 
 class _CIDPydanticAnnotation:
     @classmethod

diff --git a/tests/models/tests/test_blob_ref.py b/tests/models/tests/test_blob_ref.py
@@ -0,0 +1,31 @@
+from atproto_client.models import get_or_create
+from atproto_client.models.blob_ref import BlobRef
+from atproto_core.cid import CID
+
+
+def test_blob_ref_from_ipld() -> None:
+    plain_cid = 'bafyreidfayvfuwqa7qlnopdjiqrxzs6blmoeu4rujcjtnci5beludirz2a'
+    plain_blob_ref = {
+        'mimeType': 'text/plain',
+        'size': 0,
+        '$type': 'blob',
+        'ref': plain_cid,
+    }
+    instance = get_or_create(plain_blob_ref, BlobRef)
+
+    assert isinstance(instance, BlobRef)
+    assert instance.cid == CID.decode(plain_cid)
+
+
+def test_blob_ref_from_ipld_json() -> None:
+    plain_cid = 'bafyreidfayvfuwqa7qlnopdjiqrxzs6blmoeu4rujcjtnci5beludirz2a'
+    plain_blob_ref = {
+        'mimeType': 'text/plain',
+        'size': 0,
+        '$type': 'blob',
+        'ref': {'$link': plain_cid},
+    }
+    instance = get_or_create(plain_blob_ref, BlobRef)
+
+    assert isinstance(instance, BlobRef)
+    assert instance.cid == CID.decode(plain_cid)