Speed up model parsing, cut memory and import time

packages/atproto/__init__.py

@@ -1,42 +1,97 @@
-from atproto_client import AsyncClient, Client, Session, SessionEvent, models
-from atproto_client import utils as client_utils
-from atproto_client.request import AsyncRequest, Request
-from atproto_core.car import CAR
-from atproto_core.cid import CID, CIDType
-from atproto_core.did_doc import DidDocument
-from atproto_core.nsid import NSID
-from atproto_core.uri import AtUri
-from atproto_crypto.did import Multikey, get_did_key
-from atproto_crypto.multibase import bytes_to_multibase, multibase_to_bytes
-from atproto_crypto.verify import verify_signature
-from atproto_firehose import (
-    AsyncFirehoseSubscribeLabelsClient,
-    AsyncFirehoseSubscribeReposClient,
-    FirehoseSubscribeLabelsClient,
-    FirehoseSubscribeReposClient,
-    parse_subscribe_labels_message,
-    parse_subscribe_repos_message,
-)
-from atproto_identity.cache.in_memory_cache import AsyncDidInMemoryCache, DidInMemoryCache
-from atproto_identity.did.atproto_data import AtprotoData
-from atproto_identity.resolver import AsyncIdResolver, IdResolver
-from atproto_jetstream import (
-    AsyncJetstreamClient,
-    JetstreamClient,
-    parse_subscribe_events_message,
-)
-from atproto_jetstream import archive as jetstream_archive
-from atproto_jetstream import models as jetstream_models
-from atproto_server.auth.jwt import (
-    JwtPayload,
-    decode_jwt_payload,
-    get_jwt_payload,
-    parse_jwt,
-    validate_jwt_payload,
-    verify_jwt,
-    verify_jwt_async,
+"""The AT Protocol SDK.
+
+Names are imported on first access, so ``from atproto import Client`` does not load the firehose,
+Jetstream, identity, crypto, and server packages.
+"""
+
+import typing as t
+
+from atproto_core.lazy_exports import lazy_exports
+
+if t.TYPE_CHECKING:
+    from atproto_client import AsyncClient, Client, Session, SessionEvent, models
+    from atproto_client import utils as client_utils
+    from atproto_client.request import AsyncRequest, Request
+    from atproto_core.car import CAR
+    from atproto_core.cid import CID, CIDType
+    from atproto_core.did_doc import DidDocument
+    from atproto_core.nsid import NSID
+    from atproto_core.uri import AtUri
+    from atproto_crypto.did import Multikey, get_did_key
+    from atproto_crypto.multibase import bytes_to_multibase, multibase_to_bytes
+    from atproto_crypto.verify import verify_signature
+    from atproto_firehose import (
+        AsyncFirehoseSubscribeLabelsClient,
+        AsyncFirehoseSubscribeReposClient,
+        FirehoseSubscribeLabelsClient,
+        FirehoseSubscribeReposClient,
+        parse_subscribe_labels_message,
+        parse_subscribe_repos_message,
+    )
+    from atproto_identity.cache.in_memory_cache import AsyncDidInMemoryCache, DidInMemoryCache
+    from atproto_identity.did.atproto_data import AtprotoData
+    from atproto_identity.resolver import AsyncIdResolver, IdResolver
+    from atproto_jetstream import (
+        AsyncJetstreamClient,
+        JetstreamClient,
+        parse_subscribe_events_message,
+    )
+    from atproto_jetstream import archive as jetstream_archive
+    from atproto_jetstream import models as jetstream_models
+    from atproto_server.auth.jwt import (
+        JwtPayload,
+        decode_jwt_payload,
+        get_jwt_payload,
+        parse_jwt,
+        validate_jwt_payload,
+        verify_jwt,
+        verify_jwt_async,
+    )
+    from atproto_subscription import frames as firehose_models
+
+__getattr__, __dir__ = lazy_exports(
+    __name__,
+    attributes={
+        'atproto_client': ('AsyncClient', 'Client', 'Session', 'SessionEvent'),
+        'atproto_client.request': ('AsyncRequest', 'Request'),
+        'atproto_core.car': ('CAR',),
+        'atproto_core.cid': ('CID', 'CIDType'),
+        'atproto_core.did_doc': ('DidDocument',),
+        'atproto_core.nsid': ('NSID',),
+        'atproto_core.uri': ('AtUri',),
+        'atproto_crypto.did': ('Multikey', 'get_did_key'),
+        'atproto_crypto.multibase': ('bytes_to_multibase', 'multibase_to_bytes'),
+        'atproto_crypto.verify': ('verify_signature',),
+        'atproto_firehose': (
+            'AsyncFirehoseSubscribeLabelsClient',
+            'AsyncFirehoseSubscribeReposClient',
+            'FirehoseSubscribeLabelsClient',
+            'FirehoseSubscribeReposClient',
+            'parse_subscribe_labels_message',
+            'parse_subscribe_repos_message',
+        ),
+        'atproto_identity.cache.in_memory_cache': ('AsyncDidInMemoryCache', 'DidInMemoryCache'),
+        'atproto_identity.did.atproto_data': ('AtprotoData',),
+        'atproto_identity.resolver': ('AsyncIdResolver', 'IdResolver'),
+        'atproto_jetstream': ('AsyncJetstreamClient', 'JetstreamClient', 'parse_subscribe_events_message'),
+        'atproto_server.auth.jwt': (
+            'JwtPayload',
+            'decode_jwt_payload',
+            'get_jwt_payload',
+            'parse_jwt',
+            'validate_jwt_payload',
+            'verify_jwt',
+            'verify_jwt_async',
+        ),
+    },
+    modules={
+        'models': 'atproto_client.models',
+        'client_utils': 'atproto_client.utils',
+        'firehose_models': 'atproto_subscription.frames',
+        'jetstream_archive': 'atproto_jetstream.archive',
+        'jetstream_models': 'atproto_jetstream.models',
+    },
 )
-from atproto_subscription import frames as firehose_models
 
 __all__ = [  # noqa: RUF022
     # client

packages/atproto_client/__init__.py

@@ -1,7 +1,12 @@
-from atproto_client import models
-from atproto_client.client.async_client import AsyncClient
-from atproto_client.client.client import Client
-from atproto_client.client.session import Session, SessionEvent
+import typing as t
+
+from atproto_core.lazy_exports import lazy_exports
+
+if t.TYPE_CHECKING:
+    from atproto_client import models
+    from atproto_client.client.async_client import AsyncClient
+    from atproto_client.client.client import Client
+    from atproto_client.client.session import Session, SessionEvent
 
 __all__ = [
     'AsyncClient',
@@ -10,3 +15,12 @@
     'SessionEvent',
     'models',
 ]
+
+__getattr__, __dir__ = lazy_exports(
+    __name__,
+    attributes={
+        'atproto_client.client.async_client': ('AsyncClient',),
+        'atproto_client.client.client': ('Client',),
+        'atproto_client.client.session': ('Session', 'SessionEvent'),
+    },
+)

packages/atproto_client/models/base.py

@@ -1,50 +1,135 @@
 import os
+import threading
 import typing as t
 import warnings
 
+import typing_extensions as te
 from pydantic import BaseModel, ConfigDict, Field, alias_generators, model_validator
 
 from atproto_client.exceptions import ModelFieldNotFoundError
 
+#: Most models come in a couple of shapes. The cap bounds memory if some producer varies them a lot.
+_MAX_SHARED_FIELDS_SETS_PER_MODEL = 16
+
 
 def _alias_generator(name: str) -> str:
     return alias_generators.to_camel(name).removesuffix('_')
 
 
-class AtProtocolBase: ...
+class _SharedFieldsSet(t.Set[str]):
+    """``__pydantic_fields_set__`` shared by all models with the same set fields. Never mutated in place.
+
+    pydantic-core accepts only a :obj:`set` there, so it is a subclass that marks the set as shared.
+    Copies and pickles are plain sets.
+    """
+
+    __slots__ = ()
+
+    def __reduce__(self) -> t.Tuple[t.Any, ...]:
+        return set, (list(self),)
+
+
+#: Without the GIL, two threads unsharing one model at once would each install a copy and lose an update.
+_UNSHARE_LOCK = threading.Lock()
+
+
+def _unshare_fields_set(model: BaseModel) -> None:
+    if type(model.__pydantic_fields_set__) is not _SharedFieldsSet:
+        return
+
+    with _UNSHARE_LOCK:
+        fields_set = model.__pydantic_fields_set__
+        if type(fields_set) is _SharedFieldsSet:
+            object.__setattr__(model, '__pydantic_fields_set__', set(fields_set))
+
+
+class AtProtocolBase:
+    __slots__ = ()
 
 
 class ModelBase(BaseModel, AtProtocolBase):
     """Base class for all data classes.
 
     Provides square brackets [] notation to get attributes like in a dictionary.
+
+    Note:
+        Models with the same set fields share one ``__pydantic_fields_set__`` until one of them is modified.
+        A set per model is often the largest part of its memory.
     """
 
     model_config = ConfigDict(extra='allow', alias_generator=_alias_generator, populate_by_name=True, strict=True)
 
+    __atproto_shared_fields_sets__: t.ClassVar[t.List[_SharedFieldsSet]] = []
+
+    @classmethod
+    @te.override
+    def __pydantic_init_subclass__(cls, **kwargs: t.Any) -> None:
+        super().__pydantic_init_subclass__(**kwargs)
+        cls.__atproto_shared_fields_sets__ = []
+
     def __getitem__(self, item: str) -> t.Any:
         if hasattr(self, item):
             return getattr(self, item)
 
         raise ModelFieldNotFoundError(f"Can't find field '{item}' in the object of type {type(self)}.")
 
     @model_validator(mode='after')
-    def __alert_about_extra_fields(self) -> 'ModelBase':
-        # used for debugging purposes because gives false positives due to the way pydantic works
-
-        if self.model_extra and os.environ.get('ATPROTO_LEXICON_WARN', '0') == '1':
-            warnings.warn(
-                f'Extra fields found in the object of type {type(self)}: {self.model_extra}. '  # noqa: S608
-                f'Probably you are using the old version of SDK. '
-                f'Please update it using `pip install -U atproto`. '
-                f'In case you are working with custom lexicon ignore this warning. '
-                f'It is also possible that you are working with extended record. '
-                f'To disable this warning set `ATPROTO_LEXICON_WARN` to `0` in the environment variables.',
-                stacklevel=0,
-            )
+    def __after_validation(self) -> 'ModelBase':
+        if self.__pydantic_extra__:
+            if os.environ.get('ATPROTO_LEXICON_WARN', '0') == '1':
+                # used for debugging purposes because gives false positives due to the way pydantic works
+                warnings.warn(
+                    f'Extra fields found in the object of type {type(self)}: {self.model_extra}. '  # noqa: S608
+                    f'Probably you are using the old version of SDK. '
+                    f'Please update it using `pip install -U atproto`. '
+                    f'In case you are working with custom lexicon ignore this warning. '
+                    f'It is also possible that you are working with extended record. '
+                    f'To disable this warning set `ATPROTO_LEXICON_WARN` to `0` in the environment variables.',
+                    stacklevel=0,
+                )
+
+            # extra keys are arbitrary, so models with them are not worth sharing a set
+            return self
+
+        fields_set = self.__pydantic_fields_set__
+        shared_sets = type(self).__atproto_shared_fields_sets__
+        for shared in shared_sets:
+            # cheaper than building a hashable key
+            if shared == fields_set:
+                object.__setattr__(self, '__pydantic_fields_set__', shared)
+                return self
+
+        if len(shared_sets) < _MAX_SHARED_FIELDS_SETS_PER_MODEL:
+            shared = _SharedFieldsSet(fields_set)
+            shared_sets.append(shared)
+            object.__setattr__(self, '__pydantic_fields_set__', shared)
 
         return self
 
+    @te.override
+    def __setattr__(self, name: str, value: t.Any) -> None:
+        _unshare_fields_set(self)
+        super().__setattr__(name, value)
+
+    @te.override
+    def __setstate__(self, state: t.Dict[t.Any, t.Any]) -> None:
+        super().__setstate__(state)
+        # models pickled together come back holding one plain set, so each needs its own
+        object.__setattr__(self, '__pydantic_fields_set__', set(self.__pydantic_fields_set__))
+
+    @property
+    @te.override
+    def model_fields_set(self) -> t.Set[str]:
+        """The set of fields that have been explicitly set on this model instance."""
+        _unshare_fields_set(self)
+        return self.__pydantic_fields_set__
+
+    @property
+    @te.override
+    def __fields_set__(self) -> t.Set[str]:
+        _unshare_fields_set(self)
+        return super().__fields_set__  # pyrefly: ignore[deprecated]
+
 
 class ParamsModelBase(ModelBase):
     pass
@@ -63,7 +148,7 @@ class SugarResponseModelBase(ResponseModelBase):
 
 
 class UnknownDict(AtProtocolBase):
-    pass
+    __slots__ = ()
 
 
 class UnknownRecord(ModelBase):

packages/atproto_client/models/bytes_type.py

@@ -1,6 +1,7 @@
 import base64
 import binascii
 import typing as t
+from collections.abc import Mapping
 
 import typing_extensions as te
 from pydantic import BeforeValidator, PlainSerializer
@@ -10,7 +11,9 @@
 
 def _decode_json_bytes(value: object) -> object:
     """Decode the atproto JSON form ``{"$bytes": "<base64>"}``. Raw bytes (from CBOR) pass through."""
-    if not isinstance(value, t.Mapping) or set(value) != {_BYTES_KEY}:
+    if type(value) is bytes:
+        return value
+    if not isinstance(value, Mapping) or set(value) != {_BYTES_KEY}:
         return value
 
     encoded: object = value[_BYTES_KEY]