Use jiter instead of built-in json module to improve performance

MarshalX · Jul 26, 2024 · db72d78 · db72d78
1 parent ddc6e35
commit db72d78
Show file tree

Hide file tree

Showing 10 changed files with 56 additions and 37 deletions.
diff --git a/packages/atproto_client/models/utils.py b/packages/atproto_client/models/utils.py
@@ -1,9 +1,9 @@
-import json
 import types
 import typing as t
 
 import typing_extensions as te
 from pydantic import ValidationError
+from pydantic_core import from_json, to_json
 
 from atproto_client import models
 from atproto_client.exceptions import (
@@ -126,19 +126,19 @@ def get_model_as_dict(model: t.Union[DotDict, BlobRef, ModelBase]) -> t.Dict[str
 
 def get_model_as_json(model: t.Union[DotDict, BlobRef, ModelBase]) -> str:
     if isinstance(model, DotDict):
-        return json.dumps(get_model_as_dict(model))
+        return to_json(get_model_as_dict(model)).decode('UTF-8')
 
     return model.model_dump_json(exclude_none=True, by_alias=True)
 
 
 def is_json(json_data: t.Union[str, bytes]) -> bool:
     if isinstance(json_data, bytes):
-        json_data.decode('UTF-8')
+        json_data.decode('UTF-8', errors='ignore')
 
     try:
-        json.loads(json_data)
+        from_json(json_data)
         return True
-    except:  # noqa
+    except ValueError:
         return False
 
 

diff --git a/packages/atproto_client/request.py b/packages/atproto_client/request.py
@@ -1,9 +1,9 @@
-import json
 import typing as t
 from dataclasses import dataclass
 
 import httpx
 import typing_extensions as te
+from pydantic_core import from_json
 
 from atproto_client import exceptions
 from atproto_client.models.common import XrpcError
@@ -66,7 +66,7 @@ def _handle_response(response: httpx.Response) -> httpx.Response:
         headers=_convert_headers_to_dict(response.headers),
     )
     if response.content and is_json(response.content):
-        data: t.Dict[str, t.Any] = json.loads(response.content)
+        data: t.Dict[str, t.Any] = from_json(response.content)
         error_response.content = t.cast(XrpcError, get_or_create(data, XrpcError, strict=False))
 
     if response.status_code in {401, 403}:

diff --git a/packages/atproto_lexicon/parser.py b/packages/atproto_lexicon/parser.py
@@ -1,8 +1,9 @@
-import json
 import os
 import typing as t
 from pathlib import Path
 
+from pydantic_core import from_json
+
 from atproto_lexicon import models
 from atproto_lexicon.exceptions import LexiconParsingError
 
@@ -19,7 +20,7 @@ def lexicon_parse(data: dict, model_class: t.Optional[t.Type[L]] = models.Lexico
 def lexicon_parse_file(lexicon_path: t.Union[Path, str], *, soft_fail: bool = False) -> t.Optional[models.LexiconDoc]:
     try:
         with open(lexicon_path, encoding='UTF-8') as f:
-            plain_lexicon = json.loads(f.read())
+            plain_lexicon = from_json(f.read())
             return lexicon_parse(plain_lexicon)
     except Exception as e:  # noqa: BLE001
         if soft_fail:

diff --git a/packages/atproto_server/auth/jwt.py b/packages/atproto_server/auth/jwt.py
@@ -1,10 +1,10 @@
 import binascii
-import json
 import typing as t
 from datetime import datetime, timezone
 
 from atproto_crypto.verify import verify_signature
 from pydantic import BaseModel, ConfigDict
+from pydantic_core import from_json
 
 from atproto_server.auth.utils import base64url_decode
 from atproto_server.exceptions import (
@@ -64,14 +64,14 @@ def parse_jwt(jwt: t.Union[str, bytes]) -> t.Tuple[bytes, bytes, t.Dict[str, t.A
         raise TokenDecodeError('Invalid header padding') from e
 
     try:
-        header = json.loads(header_data)
+        header = from_json(header_data)
     except ValueError as e:
         raise TokenDecodeError(f'Invalid header string: {e}') from e
 
     if not isinstance(header, dict):
         raise TokenDecodeError('Invalid header string: must be a json object')
 
-    header = t.cast(t.Dict[str, t.Any], json.loads(header_data))  # we expect object in header
+    header = t.cast(t.Dict[str, t.Any], from_json(header_data))  # we expect an object in header
 
     try:
         payload = base64url_decode(payload_segment)
@@ -96,7 +96,7 @@ def decode_jwt_payload(payload: t.Union[str, bytes]) -> JwtPayload:
         :obj:`JwtPayload`: The decoded payload of the given JWT.
     """
     try:
-        plain_payload = json.loads(payload)
+        plain_payload = from_json(payload)
     except ValueError as e:
         raise TokenDecodeError(f'Invalid payload string: {e}') from e
     if not isinstance(plain_payload, dict):

diff --git a/tests/test_atproto_client/models/fetch_test_data.py b/tests/test_atproto_client/models/fetch_test_data.py
@@ -1,10 +1,10 @@
-import json
 import logging
 import os
 import typing as t
 
 from atproto_client import Client
 from atproto_client.request import Request, Response
+from pydantic_core import to_json
 
 if t.TYPE_CHECKING:
     from atproto_client.models.common import XrpcError
@@ -73,7 +73,7 @@ def get_unique_filename(name: str) -> str:
 
 
 def get_pretty_json(data: dict) -> str:
-    return json.dumps(data, indent=4)
+    return to_json(data, indent=4).decode('UTF-8')
 
 
 def save_response(name: str) -> None:

diff --git a/tests/test_atproto_client/models/test_data/feed_record.json b/tests/test_atproto_client/models/test_data/feed_record.json
@@ -2,7 +2,6 @@
     "uri": "at://did:plc:s6jnht6koorxz7trghirytmf/app.bsky.feed.generator/atproto",
     "cid": "bafyreifiexhek65jxj3ucz6y6zstj45rmtigybzygjkg6lretyqgtge5ai",
     "value": {
-        "did": "did:web:feed.atproto.blue",
         "$type": "app.bsky.feed.generator",
         "avatar": {
             "$type": "blob",
@@ -14,6 +13,7 @@
         },
         "createdAt": "2023-07-20T10:17:40.298101",
         "description": "Posts related to the protocol. Powered by The AT Protocol SDK for Python",
+        "did": "did:web:feed.atproto.blue",
         "displayName": "AT Protocol"
     }
 }
diff --git a/tests/test_atproto_client/models/test_data/get_follows.json b/tests/test_atproto_client/models/test_data/get_follows.json
@@ -7,7 +7,7 @@
             "avatar": "https://cdn.bsky.app/img/avatar/plain/did:plc:vpkhqolt662uhesyj6nxm7ys/bafkreietenc2aywhzxdyvalgffxl35tvxri6wv2iwryfdcdoecdfg553fy@jpeg",
             "associated": {
                 "chat": {
-                    "allowIncoming": "following"
+                    "allowIncoming": "all"
                 }
             },
             "viewer": {
@@ -16,17 +16,18 @@
                 "following": "at://did:plc:kvwvcn5iqfooopmyzvb4qzba/app.bsky.graph.follow/3jueqtsliw22n"
             },
             "labels": [],
-            "description": "Technical advisor to @bluesky, first engineer at Protocol Labs. Wizard Utopian",
-            "indexedAt": "2024-05-14T15:42:22.889Z"
+            "createdAt": "2022-11-17T01:04:43.624Z",
+            "description": "Technical advisor to @bluesky, first engineer at Protocol Labs. Wizard Utopian.",
+            "indexedAt": "2024-06-07T03:19:53.642Z"
         },
         {
             "did": "did:plc:l3rouwludahu3ui3bt66mfvj",
             "handle": "divy.zone",
-            "displayName": "devin ivy \ud83d\udc0b",
+            "displayName": "devin ivy 🐋",
             "avatar": "https://cdn.bsky.app/img/avatar/plain/did:plc:l3rouwludahu3ui3bt66mfvj/bafkreicg6y3mlr3eszmbjm3swyuncwf4ruzohcsvtcbrjzyhkwibtx7nyy@jpeg",
             "associated": {
                 "chat": {
-                    "allowIncoming": "following"
+                    "allowIncoming": "all"
                 }
             },
             "viewer": {
@@ -36,6 +37,7 @@
                     "cid": "bafyreigk3kmjipz5emav5vqmvdtivwz757tfpg5lnsbfoscnqx7wpjjime",
                     "name": "test mute list",
                     "purpose": "app.bsky.graph.defs#modlist",
+                    "listItemCount": 1,
                     "indexedAt": "2023-08-28T10:08:27.442Z",
                     "labels": [],
                     "viewer": {
@@ -46,13 +48,14 @@
                 "following": "at://did:plc:kvwvcn5iqfooopmyzvb4qzba/app.bsky.graph.follow/3jueqt6dbqs2g"
             },
             "labels": [],
-            "description": "\ud83c\udf00 bluesky team",
-            "indexedAt": "2024-03-08T04:03:32.618Z"
+            "createdAt": "2022-11-17T00:39:19.084Z",
+            "description": "🌀 bluesky team",
+            "indexedAt": "2024-06-14T20:22:02.642Z"
         },
         {
             "did": "did:plc:oky5czdrnfjpqslsw2a5iclo",
             "handle": "jay.bsky.team",
-            "displayName": "Jay \ud83e\udd8b",
+            "displayName": "Jay 🦋",
             "avatar": "https://cdn.bsky.app/img/avatar/plain/did:plc:oky5czdrnfjpqslsw2a5iclo/bafkreihidru2xruxdxlvvcixc7lbgoudzicjbrdgacdhdhxyfw4yut4nfq@jpeg",
             "associated": {
                 "chat": {
@@ -65,7 +68,8 @@
                 "following": "at://did:plc:kvwvcn5iqfooopmyzvb4qzba/app.bsky.graph.follow/3judl7ak7gp2f"
             },
             "labels": [],
-            "description": "CEO of Bluesky, steward of AT Protocol. \n\nLet\u2019s build a federated republic, starting with this server. \ud83c\udf31 \ud83e\udeb4 \ud83c\udf33 ",
+            "createdAt": "2022-11-17T06:31:40.296Z",
+            "description": "CEO of Bluesky, steward of AT Protocol. \n\nLet’s build a federated republic, starting with this server. 🌱 🪴 🌳 ",
             "indexedAt": "2024-02-06T22:21:45.352Z"
         },
         {
@@ -84,21 +88,28 @@
                 "following": "at://did:plc:kvwvcn5iqfooopmyzvb4qzba/app.bsky.graph.follow/3judkza2vrb2y"
             },
             "labels": [],
-            "description": "Official Bluesky account (check domain\ud83d\udc46)\n\nFollow for updates and announcements",
+            "createdAt": "2023-04-12T04:53:57.057Z",
+            "description": "Official Bluesky account (check domain👆)\n\nFollow for updates and announcements",
             "indexedAt": "2024-01-25T23:46:28.776Z"
         },
         {
             "did": "did:plc:m2jwplpernhxkzbo4ev5ljwj",
             "handle": "vercel.com",
             "displayName": "Vercel",
             "avatar": "https://cdn.bsky.app/img/avatar/plain/did:plc:m2jwplpernhxkzbo4ev5ljwj/bafkreicebob2yf5lv6yg72luzv5qwsr6ob65j6oc3jciyowqkfiz736oqu@jpeg",
+            "associated": {
+                "chat": {
+                    "allowIncoming": "following"
+                }
+            },
             "viewer": {
                 "muted": false,
                 "blockedBy": false,
                 "following": "at://did:plc:kvwvcn5iqfooopmyzvb4qzba/app.bsky.graph.follow/3judkwijszc25"
             },
             "labels": [],
-            "description": "Vercel\u2019s frontend cloud gives developers the frameworks, workflows, and infrastructure to build a faster, more personalized web. Creators of @nextjs.org.",
+            "createdAt": "2023-04-25T00:08:45.850Z",
+            "description": "Vercel’s frontend cloud gives developers the frameworks, workflows, and infrastructure to build a faster, more personalized web. Creators of @nextjs.org.",
             "indexedAt": "2024-02-16T21:57:28.740Z"
         },
         {
@@ -118,7 +129,8 @@
                 "followedBy": "at://did:plc:s6jnht6koorxz7trghirytmf/app.bsky.graph.follow/3jucc25a7qs2k"
             },
             "labels": [],
-            "description": "Software Engineer\n\n\ud83d\udc0d The AT Protocol SDK for Python: https://atproto.blue/\n\ud83c\udf7f Custom Feed in Python: https://github.com/MarshalX/bluesky-feed-generator\n\ud83c\udfce\ufe0f Fast DAG-CBOR decoder for Python: https://github.com/MarshalX/python-libipld\n\nhttps://marshal.dev",
+            "createdAt": "2023-04-12T11:14:00.501Z",
+            "description": "Software Engineer\n\n🐍 The AT Protocol SDK for Python: https://atproto.blue/\n🍿 Custom Feed in Python: https://github.com/MarshalX/bluesky-feed-generator\n🏎️ Fast DAG-CBOR decoder for Python: https://github.com/MarshalX/python-libipld\n\nhttps://marshal.dev",
             "indexedAt": "2024-01-26T00:15:07.447Z"
         }
     ],
@@ -137,7 +149,8 @@
             "blockedBy": false
         },
         "labels": [],
-        "description": "account for tests",
-        "indexedAt": "2024-05-22T21:04:02.588Z"
+        "createdAt": "2023-04-26T19:05:34.249Z",
+        "description": "account for tests\n\nAuthor: @marshal.dev\nGitHub: https://github.com/MarshalX/atproto\nWebsite: https://atproto.blue",
+        "indexedAt": "2024-05-22T21:19:48.088Z"
     }
 }
diff --git a/tests/test_atproto_client/models/test_data/thread_view_post_with_embed_media.json b/tests/test_atproto_client/models/test_data/thread_view_post_with_embed_media.json
@@ -18,7 +18,8 @@
                     "muted": false,
                     "blockedBy": false
                 },
-                "labels": []
+                "labels": [],
+                "createdAt": "2023-04-26T19:05:34.249Z"
             },
             "record": {
                 "$type": "app.bsky.feed.post",
@@ -85,7 +86,8 @@
                                 "muted": false,
                                 "blockedBy": false
                             },
-                            "labels": []
+                            "labels": [],
+                            "createdAt": "2023-04-26T19:05:34.249Z"
                         },
                         "value": {
                             "$type": "app.bsky.feed.post",
@@ -108,7 +110,9 @@
             "repostCount": 0,
             "likeCount": 0,
             "indexedAt": "2023-09-28T12:49:36.735Z",
-            "viewer": {},
+            "viewer": {
+                "threadMuted": false
+            },
             "labels": []
         },
         "replies": []

diff --git a/tests/test_atproto_client/models/tests/utils.py b/tests/test_atproto_client/models/tests/utils.py
@@ -1,9 +1,10 @@
-import json
 import os
 
+from pydantic_core import from_json
+
 TEST_DATA_PATH = os.path.join(os.path.dirname(__file__), '..', 'test_data')
 
 
 def load_data_from_file(test_name: str) -> dict:
     with open(os.path.join(TEST_DATA_PATH, f'{test_name}.json'), 'r') as f:
-        return json.load(f)
+        return from_json(f.read())
diff --git a/tests/test_atproto_crypto/test_verify.py b/tests/test_atproto_crypto/test_verify.py
@@ -1,17 +1,17 @@
 import base64
-import json
 import os
 
 import pytest
 from atproto_crypto.verify import verify_signature
+from pydantic_core import from_json
 
 # Ref: https://github.com/bluesky-social/atproto/blob/main/interop-test-files/crypto/signature-fixtures.json
 _FIXTURES_FILE_PATH = os.path.join(os.path.dirname(__file__), 'signature-fixtures.json')
 
 
 def _load_test_cases() -> list:
     with open(_FIXTURES_FILE_PATH, encoding='UTF-8') as file:
-        return json.load(file)
+        return from_json(file.read())
 
 
 def _fix_base64_padding(data: str) -> str: