Fix parsing of custom (extended) records; add dot notation for dictio…

…naries (#106)

atproto/codegen/models/generator.py

@@ -153,6 +153,7 @@ def _get_ref_union_typehint(nsid: NSID, field_type_def, *, optional: bool) -> st
     # maybe it's for the records that have custom fields... idk
     # ref: https://github.com/bluesky-social/atproto/blob/b01e47b61730d05a780f7a42667b91ccaa192e8e/packages/lex-cli/src/codegen/lex-gen.ts#L325
     # grep by "{$type: string; [k: string]: unknown}" string
+    # TODO(MarshalX): use 'base.UnknownDict' and convert to DotDict
     def_names.append('t.Dict[str, t.Any]')
 
     def_names = ', '.join([f"'{name}'" for name in def_names])

atproto/xrpc_client/models/utils.py

@@ -29,7 +29,7 @@
 def _unknown_type_hook(data: dict) -> t.Union[UnknownRecordType, DotDict]:
     if '$type' in data:
         # $type used for inner Record types
-        return get_or_create_model(data, RECORD_TYPE_TO_MODEL_CLASS[data.pop('$type')])
+        return get_or_create(data, strict=False)
     # any another unknown (not described by lexicon) type
     return DotDict(data)
 
@@ -51,7 +51,7 @@ def _decode_cid_hook(ref: t.Union[CID, str]) -> CID:
 
 def get_or_create(
     model_data: ModelData, model: t.Optional[t.Type[M]] = None, *, strict: bool = True
-) -> t.Optional[t.Union[M, dict]]:
+) -> t.Optional[t.Union[M, UnknownRecordType, DotDict]]:
     """Get model instance from raw data.
 
     Note:
@@ -80,17 +80,19 @@ def get_or_create(
         if strict:
             raise e
 
-        return model_data
+        return DotDict(model_data)
 
 
-def _get_or_create(model_data: ModelData, model: t.Type[M], *, strict: bool) -> t.Optional[t.Union[M, dict]]:
+def _get_or_create(
+    model_data: ModelData, model: t.Type[M], *, strict: bool
+) -> t.Optional[t.Union[M, UnknownRecordType, DotDict]]:
     if model_data is None:
         return None
 
     if model is None:
         # resolve model by $type and try to parse
         # resolves only Records
-        record_type = model_data.pop('$type')
+        record_type = model_data.pop('$type', None)
         if not record_type or record_type not in RECORD_TYPE_TO_MODEL_CLASS:
             return None
 

atproto/xrpc_client/request.py

@@ -93,7 +93,7 @@ def __init__(self) -> None:
         self._client = httpx.Client()
 
     def _send_request(self, method: str, url: str, **kwargs) -> httpx.Response:
-        headers = self.get_headers(kwargs.pop('headers'))
+        headers = self.get_headers(kwargs.pop('headers', None))
 
         try:
             response = self._client.request(method=method, url=url, headers=headers, **kwargs)
@@ -120,7 +120,7 @@ def __init__(self) -> None:
         self._client = httpx.AsyncClient()
 
     async def _send_request(self, method: str, url: str, **kwargs) -> httpx.Response:
-        headers = self.get_headers(kwargs.pop('headers'))
+        headers = self.get_headers(kwargs.pop('headers', None))
 
         try:
             response = await self._client.request(method=method, url=url, headers=headers, **kwargs)

test.py

@@ -46,13 +46,23 @@ def sync_main():
     client = Client()
     client.login(os.environ['USERNAME'], os.environ['PASSWORD'])
 
-    # r = client.send_post('regress test')
-    # print(r)
+    lexicon_correct_record = client.com.atproto.repo.get_record(
+        {'collection': 'app.bsky.feed.post', 'repo': 'test.marshal.dev', 'rkey': '3k2yihcrp6f2c'}
+    )
+    print(lexicon_correct_record.value.text)
+    print(type(lexicon_correct_record.value))
+    extended_record = client.com.atproto.repo.get_record(
+        {'collection': 'app.bsky.feed.post', 'repo': 'test.marshal.dev', 'rkey': '3k2yinh52ne2x'}
+    )
+    print(extended_record.value.text)
+    print(extended_record.value.lol)  # custom (out of lexicon) attribute
+    print(type(extended_record.value))
 
     did_doc = client.com.atproto.repo.describe_repo({'repo': 'did:plc:ze3uieyyns7prike7itbdjiy'}).didDoc
     print(did_doc.service)
     print(did_doc['service'])
     print(did_doc['@context'])
+    print(type(did_doc))
 
     atproto_feed = client.com.atproto.repo.get_record(
         {'collection': ids.AppBskyFeedGenerator, 'repo': 'marshal.dev', 'rkey': 'atproto'}
@@ -116,7 +126,7 @@ async def main():
     # with open('cat2.png', 'rb') as f:
     #     cat_data = f.read()
 
-    # await async_client.send_image('Cat looking for a Async Python', cat_data, 'async cat alt')
+    # await async_client.send_image('Cat looking for an Async Python', cat_data, 'async cat alt')
 
     # resolve = await async_client.com.atproto.identity.resolve_handle(
     #     models.ComAtprotoIdentityResolveHandle.Params(profile.handle)
@@ -171,11 +181,9 @@ def _get_ops_by_type(commit: models.ComAtprotoSyncSubscribeRepos.Commit) -> dict
 
 
 def _custom_feed_firehose():
-    client = FirehoseSubscribeReposClient({'cursor': 93278360})
     client = FirehoseSubscribeReposClient()
 
     def on_message_handler(message: 'MessageFrame') -> None:
-        # return
         commit = parse_subscribe_repos_message(message)
         if not isinstance(commit, models.ComAtprotoSyncSubscribeRepos.Commit):
             return