From 33529b0002537f1f00965ae180639bf1663e8ed8 Mon Sep 17 00:00:00 2001 From: "Ilya (Marshal)" Date: Thu, 22 Oct 2026 21:52:47 +0200 Subject: [PATCH] Speed up model parsing, cut memory and import time --- packages/atproto/__init__.py | 131 +++++++++++++----- packages/atproto_client/__init__.py | 22 ++- packages/atproto_client/models/base.py | 115 +++++++++++++-- packages/atproto_client/models/bytes_type.py | 5 +- packages/atproto_client/models/dot_dict.py | 56 +++++--- .../atproto_client/models/models_loader.py | 3 +- .../atproto_client/models/string_formats.py | 7 +- .../atproto_client/models/unknown_union.py | 13 +- packages/atproto_client/models/utils.py | 11 +- packages/atproto_core/cid/cid.py | 13 +- packages/atproto_core/lazy_exports.py | 50 +++++++ packages/atproto_crypto/did.py | 20 ++- packages/atproto_crypto/verify.py | 4 +- packages/atproto_subscription/frames.py | 8 +- tests/test_atproto/test_lazy_imports.py | 85 ++++++++++++ tests/test_atproto/test_models.py | 5 + .../models/tests/test_dot_dict.py | 39 ++++++ .../models/tests/test_shared_fields_set.py | 87 ++++++++++++ tests/test_atproto_core/test_cid.py | 20 +++ 19 files changed, 583 insertions(+), 111 deletions(-) create mode 100644 packages/atproto_core/lazy_exports.py create mode 100644 tests/test_atproto/test_lazy_imports.py create mode 100644 tests/test_atproto_client/models/tests/test_shared_fields_set.py diff --git a/packages/atproto/__init__.py b/packages/atproto/__init__.py index 91b1df83d..e51a4666a 100644 --- a/packages/atproto/__init__.py +++ b/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 diff --git a/packages/atproto_client/__init__.py b/packages/atproto_client/__init__.py index 9c8dd1467..5d7df7fb2 100644 --- a/packages/atproto_client/__init__.py +++ b/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'), + }, +) diff --git a/packages/atproto_client/models/base.py b/packages/atproto_client/models/base.py index e22c6ff31..4525d8b89 100644 --- a/packages/atproto_client/models/base.py +++ b/packages/atproto_client/models/base.py @@ -1,27 +1,72 @@ 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) @@ -29,22 +74,62 @@ def __getitem__(self, item: str) -> t.Any: 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): diff --git a/packages/atproto_client/models/bytes_type.py b/packages/atproto_client/models/bytes_type.py index baba366cb..7bfd84927 100644 --- a/packages/atproto_client/models/bytes_type.py +++ b/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": ""}``. 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] diff --git a/packages/atproto_client/models/dot_dict.py b/packages/atproto_client/models/dot_dict.py index 55458bcbc..73297d6fe 100644 --- a/packages/atproto_client/models/dot_dict.py +++ b/packages/atproto_client/models/dot_dict.py @@ -22,7 +22,8 @@ def _convert_snake_case_to_camel_case(string: str) -> str: def _is_snake_case(string: str) -> bool: - return string == _convert_camel_case_to_snake_case(string) + # same as `string == _convert_camel_case_to_snake_case(string)`, which only changes strings with uppercase + return string == string.lower() def _convert_to_opposite_case(string: str) -> str: @@ -62,14 +63,16 @@ class DotDict(UnknownDict): >>> assert model.to_dict() == test_data """ + __slots__ = ('__weakref__', '_data') + + _data: t.Dict[str, t.Any] + def __init__(self, data: t.Dict[str, t.Any]) -> None: - self._data = deepcopy(data) - for k, v in self._data.items(): - self.__setitem__(k, v) + self._data = {key: _deep_copy(value, DotDict) for key, value in data.items()} def to_dict(self) -> t.Dict[str, t.Any]: """Unwrap DotDict to Python built-in dict.""" - return deepcopy(self._data) + return _to_dict(self._data) def __getitem__(self, item: str) -> t.Any: value = self._data.get(item) @@ -82,14 +85,14 @@ def __getitem__(self, item: str) -> t.Any: def __setitem__(self, key: str, value: t.Any) -> None: if key == '_data': - super().__setattr__(key, value) + object.__setattr__(self, key, value) return # we store the field in case that was firstly meet to not create duplicates if key not in self._data and _is_snake_case(key): key = _convert_snake_case_to_camel_case(key) - self._data.__setitem__(key, DotDict.__convert(value)) + self._data[key] = _deep_copy(value, DotDict) __setattr__ = __setitem__ @@ -131,17 +134,36 @@ def __reduce_ex__(self, protocol: te.SupportsIndex) -> t.Union[str, t.Tuple[t.An def __reduce__(self) -> t.Union[str, t.Tuple[t.Any, ...]]: return self._data.__reduce__() - @staticmethod - def __convert(obj: t.Any) -> t.Any: - if isinstance(obj, dict): - return DotDict(t.cast('t.Dict[str, t.Any]', obj)) - if isinstance(obj, list): - return [DotDict.__convert(v) for v in t.cast('t.List[t.Any]', obj)] - if isinstance(obj, set): - return {DotDict.__convert(v) for v in t.cast('t.Set[t.Any]', obj)} - if isinstance(obj, tuple): - return tuple(DotDict.__convert(v) for v in t.cast('t.Tuple[t.Any]', obj)) + +#: Leaves that a deep copy returns as is. +_ATOMIC_TYPES = frozenset({str, int, float, bool, bytes, type(None)}) + + +def _deep_copy(obj: t.Any, copy_dict: t.Callable[[t.Dict[str, t.Any]], t.Any]) -> t.Any: + """Deep copy ``obj``, building every dict, a :obj:`DotDict`'s included, with ``copy_dict``.""" + obj_type = type(obj) + if obj_type in _ATOMIC_TYPES: return obj + # exact types first: they are the common case and skip the isinstance() checks + if obj_type is dict: + return copy_dict(obj) + if obj_type is list: + return [_deep_copy(v, copy_dict) for v in obj] + if isinstance(obj, DotDict): + return copy_dict(obj._data) + if isinstance(obj, dict): + return copy_dict(t.cast('t.Dict[str, t.Any]', obj)) + if isinstance(obj, list): + return [_deep_copy(v, copy_dict) for v in t.cast('t.List[t.Any]', obj)] + if isinstance(obj, set): + return {_deep_copy(v, copy_dict) for v in t.cast('t.Set[t.Any]', obj)} + if isinstance(obj, tuple): + return tuple(_deep_copy(v, copy_dict) for v in t.cast('t.Tuple[t.Any]', obj)) + return deepcopy(obj) + + +def _to_dict(data: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]: + return {key: _deep_copy(value, _to_dict) for key, value in data.items()} class _DotDictPydanticAnnotation: diff --git a/packages/atproto_client/models/models_loader.py b/packages/atproto_client/models/models_loader.py index 3af604793..9a0f4338d 100644 --- a/packages/atproto_client/models/models_loader.py +++ b/packages/atproto_client/models/models_loader.py @@ -97,7 +97,8 @@ def make_lazy_accessors( def __getattr__(name: str) -> t.Any: package = importlib.import_module(package_name) - nsid = getattr(package.__dict__.get('ids'), name, None) + # aliases never start with an underscore, while dunders like `__module__` resolve on `ids` too + nsid = None if name.startswith('_') else getattr(package.__dict__.get('ids'), name, None) if nsid is not None: module = importlib.import_module(_module_path_from_nsid(package_name, nsid)) prepare_model_module(module, package) diff --git a/packages/atproto_client/models/string_formats.py b/packages/atproto_client/models/string_formats.py index c1029366e..d0f40bd70 100644 --- a/packages/atproto_client/models/string_formats.py +++ b/packages/atproto_client/models/string_formats.py @@ -84,9 +84,10 @@ def __init__(self, validate_fn: Callable[..., str]) -> None: self._strict_fn: Callable[..., str] = getattr(validate_fn, '__wrapped__', validate_fn) def __call__(self, v: str, info: ValidationInfo) -> str: - if _is_strict(info): - return self._strict_fn(v, info) - return v + # runs for every string format field of every model, so the common dict context skips the call + context: object = info.context + is_strict = context.get(_OPT_IN_KEY) if type(context) is dict else _is_strict(info) + return self._strict_fn(v, info) if is_strict else v def __str__(self) -> str: func_str = f':func:`string_formats.{self.validate_fn.__name__}`' diff --git a/packages/atproto_client/models/unknown_union.py b/packages/atproto_client/models/unknown_union.py index 6c314f0f0..5869b6b3e 100644 --- a/packages/atproto_client/models/unknown_union.py +++ b/packages/atproto_client/models/unknown_union.py @@ -83,16 +83,21 @@ def __get_pydantic_core_schema__(cls, source_type: t.Any, handler: GetCoreSchema def validate( value: t.Any, validator: core_schema.ValidatorFunctionWrapHandler, info: core_schema.ValidationInfo ) -> t.Any: - if isinstance(value, BaseModel): - return value + data = value + # a dict is the common case, so it skips the slower ABC checks + if type(value) is not dict: + if isinstance(value, BaseModel): + return value + if isinstance(value, DotDict): + data = value.to_dict() - data = value.to_dict() if isinstance(value, DotDict) else value if isinstance(data, dict): record_type = data.get('$type') model = resolve_record_type(record_type) if isinstance(record_type, str) else None if model is not None: try: - return model.model_validate(data, context=info.context) + # skips the model_validate() wrapper, which only forwards the arguments + return model.__pydantic_validator__.validate_python(data, context=info.context) except ValidationError: # a record the models know of but cannot describe still degrades to DotDict pass diff --git a/packages/atproto_client/models/utils.py b/packages/atproto_client/models/utils.py index 51a9a55c4..9779d7961 100644 --- a/packages/atproto_client/models/utils.py +++ b/packages/atproto_client/models/utils.py @@ -138,7 +138,7 @@ def get_or_create( :obj:`atproto.xrpc_client.models.dot_dict.DotDict` if `strict` is disabled. """ try: - model_instance = _get_or_create(model_data, model, strict=strict, strict_string_format=strict_string_format) + model_instance = _get_or_create(model_data, model, strict_string_format=strict_string_format) if strict and model is not None and model_instance is not None and not isinstance(model_instance, model): raise ModelError(f"Can't properly parse model of type {model}") @@ -151,7 +151,7 @@ def get_or_create( def _get_or_create( - model_data: ModelData[M], model: t.Optional[t.Type[M]], *, strict: bool, strict_string_format: bool + model_data: ModelData[M], model: t.Optional[t.Type[M]], *, strict_string_format: bool ) -> t.Optional[t.Union[M, UnknownRecordType, DotDict]]: if model_data is None: return None @@ -166,12 +166,7 @@ def _get_or_create( if record_model is None: return None - return get_or_create( - model_data, - record_model, - strict=strict, - strict_string_format=strict_string_format, - ) + model = t.cast('t.Type[M]', record_model) if isinstance(model_data, model): return model_data diff --git a/packages/atproto_core/cid/cid.py b/packages/atproto_core/cid/cid.py index 379732ab3..e8fe50461 100644 --- a/packages/atproto_core/cid/cid.py +++ b/packages/atproto_core/cid/cid.py @@ -41,19 +41,10 @@ def __init__( @classmethod def decode(cls, value: t.Union[str, bytes]) -> 'CID': + # validates the CID; the multihash is rarely read, so it is decoded again on first access cid = libipld.decode_cid(value) - multihash = Multihash( - code=cid['hash']['code'], - size=cid['hash']['size'], - digest=cid['hash']['digest'], - ) - - instance = cls( - version=cid['version'], - codec=cid['codec'], - hash=multihash, - ) + instance = cls(version=cid['version'], codec=cid['codec']) if isinstance(value, str): instance._stringified_form = value diff --git a/packages/atproto_core/lazy_exports.py b/packages/atproto_core/lazy_exports.py new file mode 100644 index 000000000..3d8d13690 --- /dev/null +++ b/packages/atproto_core/lazy_exports.py @@ -0,0 +1,50 @@ +import importlib +import sys +import typing as t + + +def lazy_exports( + package_name: str, + attributes: t.Mapping[str, t.Iterable[str]], + modules: t.Optional[t.Mapping[str, str]] = None, +) -> t.Tuple[t.Callable[[str], t.Any], t.Callable[[], t.List[str]]]: + """Build ``__getattr__`` and ``__dir__`` for a package that imports its exports on first access. + + Unknown names fall back to the package's submodules. Resolved values are cached on the package. + + Args: + package_name: The package these belong to. + attributes: Module to the names it exports under the same name. + modules: Name to the module exported as is. + """ + sources: t.Dict[str, t.Tuple[str, t.Optional[str]]] = { + name: (module_name, name) for module_name, names in attributes.items() for name in names + } + for name, module_name in (modules or {}).items(): + sources[name] = (module_name, None) + + def __getattr__(name: str) -> t.Any: + source = sources.get(name) + if source is not None: + module_name, attribute = source + value = importlib.import_module(module_name) + if attribute is not None: + value = getattr(value, attribute) + + vars(sys.modules[package_name])[name] = value + return value + + # dunders are import machinery probing the package, never a submodule + if not name.startswith('_'): + try: + return importlib.import_module(f'{package_name}.{name}') + except ModuleNotFoundError as e: + if e.name != f'{package_name}.{name}': + raise + + raise AttributeError(f'module {package_name!r} has no attribute {name!r}') + + def __dir__() -> t.List[str]: + return sorted({*vars(sys.modules[package_name]), *sources}) + + return __getattr__, __dir__ diff --git a/packages/atproto_crypto/did.py b/packages/atproto_crypto/did.py index 8581704f9..a070c8ba6 100644 --- a/packages/atproto_crypto/did.py +++ b/packages/atproto_crypto/did.py @@ -1,8 +1,6 @@ import typing as t from dataclasses import dataclass -from atproto_crypto.algs.p256 import P256 -from atproto_crypto.algs.secp256k1 import Secp256k1 from atproto_crypto.consts import ( BASE58_MULTIBASE_PREFIX, DID_KEY_PREFIX, @@ -14,6 +12,9 @@ from atproto_crypto.exceptions import IncorrectDidKeyPrefixError, IncorrectMultikeyPrefixError, UnsupportedKeyTypeError from atproto_crypto.multibase import bytes_to_multibase, multibase_to_bytes +if t.TYPE_CHECKING: + from atproto_crypto.algs.base_alg import AlgBase + @dataclass class Multikey: @@ -41,6 +42,13 @@ def to_str(self) -> str: return format_multikey(self.jwt_alg, self.key_bytes) +def _get_algorithm(jwt_alg: str) -> 'AlgBase': + # imported on use: `cryptography` is slow to import, and parsing DID documents does not need it + from atproto_crypto.algs import ALGORITHM_TO_CLASS + + return ALGORITHM_TO_CLASS[jwt_alg]() + + def get_multikey_alg(multikey: str) -> str: """Get JWT alg for multikey. @@ -82,11 +90,11 @@ def parse_multikey(multikey: str) -> Multikey: if prefixed_bytes.startswith(P256_DID_PREFIX): jwt_alg = P256_JWT_ALG compressed_key_bytes = prefixed_bytes[len(P256_DID_PREFIX) :] - key_bytes = P256().decompress_pubkey(compressed_key_bytes) + key_bytes = _get_algorithm(jwt_alg).decompress_pubkey(compressed_key_bytes) elif prefixed_bytes.startswith(SECP256K1_DID_PREFIX): jwt_alg = SECP256K1_JWT_ALG compressed_key_bytes = prefixed_bytes[len(SECP256K1_DID_PREFIX) :] - key_bytes = Secp256k1().decompress_pubkey(compressed_key_bytes) + key_bytes = _get_algorithm(jwt_alg).decompress_pubkey(compressed_key_bytes) else: raise UnsupportedKeyTypeError('Unsupported key type') @@ -110,10 +118,10 @@ def format_multikey(jwt_alg: str, key: bytes) -> str: """ if jwt_alg == P256_JWT_ALG: prefix = P256_DID_PREFIX - compressed_key_bytes = P256().compress_pubkey(key) + compressed_key_bytes = _get_algorithm(jwt_alg).compress_pubkey(key) elif jwt_alg == SECP256K1_JWT_ALG: prefix = SECP256K1_DID_PREFIX - compressed_key_bytes = Secp256k1().compress_pubkey(key) + compressed_key_bytes = _get_algorithm(jwt_alg).compress_pubkey(key) else: raise UnsupportedKeyTypeError('Unsupported key type') diff --git a/packages/atproto_crypto/verify.py b/packages/atproto_crypto/verify.py index 0110587e0..8362e625e 100644 --- a/packages/atproto_crypto/verify.py +++ b/packages/atproto_crypto/verify.py @@ -1,4 +1,3 @@ -from atproto_crypto.algs import ALGORITHM_TO_CLASS from atproto_crypto.did import parse_did_key from atproto_crypto.exceptions import UnsupportedSignatureAlgorithmError @@ -14,6 +13,9 @@ def verify_signature(did_key: str, signing_input: bytes, signature: bytes) -> bo Returns: bool: True if signature is valid, False otherwise. """ + # imported on use: `cryptography` is slow to import, and decoding JWTs does not need it + from atproto_crypto.algs import ALGORITHM_TO_CLASS + parsed_did_key = parse_did_key(did_key) if parsed_did_key.jwt_alg not in ALGORITHM_TO_CLASS: raise UnsupportedSignatureAlgorithmError('Unsupported signature alg') diff --git a/packages/atproto_subscription/frames.py b/packages/atproto_subscription/frames.py index 6a41a26f8..111963fa6 100644 --- a/packages/atproto_subscription/frames.py +++ b/packages/atproto_subscription/frames.py @@ -62,13 +62,17 @@ class ErrorFrameBody: message: Optional[str] = None #: Description of the error. +#: Avoids an `Enum` call per frame. +_FRAME_TYPE_BY_OP: dict[int, FrameType] = {frame_type.value: frame_type for frame_type in FrameType} + + def parse_frame_header(raw_header: dict[str, Any]) -> FrameHeader: try: header_op = int(raw_header.get('op', 0)) - if not FrameType.has_value(header_op): + frame_type = _FRAME_TYPE_BY_OP.get(header_op) + if frame_type is None: raise FrameDecodingError('Invalid frame type') - frame_type = FrameType(header_op) if frame_type is FrameType.MESSAGE: return MessageFrameHeader(op=frame_type, t=raw_header.get('t')) return ErrorFrameHeader(op=frame_type) diff --git a/tests/test_atproto/test_lazy_imports.py b/tests/test_atproto/test_lazy_imports.py new file mode 100644 index 000000000..91b451fcf --- /dev/null +++ b/tests/test_atproto/test_lazy_imports.py @@ -0,0 +1,85 @@ +import subprocess +import sys +import typing as t + +import atproto +import pytest + + +def _imported_modules(statement: str) -> t.Set[str]: + code = f'import sys\n{statement}\nprint("\\n".join(sys.modules))' + output = subprocess.run( # noqa: S603 + [sys.executable, '-c', code], capture_output=True, text=True, check=True + ).stdout + + return set(output.split()) + + +def test_every_public_name_resolves() -> None: + for name in atproto.__all__: + assert getattr(atproto, name) is not None + assert name in dir(atproto) + + +def test_unknown_name_raises_attribute_error() -> None: + with pytest.raises(AttributeError): + _ = atproto.DoesNotExist + + +def test_importing_the_package_loads_nothing_else() -> None: + modules = _imported_modules('import atproto') + + assert {module for module in modules if module.startswith('atproto_')} == { + 'atproto_core', + 'atproto_core.lazy_exports', + } + + +def test_client_does_not_load_streams_or_crypto() -> None: + modules = _imported_modules('from atproto import Client') + + assert 'atproto_firehose' not in modules + assert 'atproto_jetstream' not in modules + assert 'cryptography' not in modules + + +def test_firehose_does_not_load_the_http_client() -> None: + modules = _imported_modules('from atproto import FirehoseSubscribeReposClient') + + assert 'httpx' not in modules + assert 'atproto_client.client.client' not in modules + + +def test_client_package_exposes_its_submodules() -> None: + code = 'import atproto_client\nprint(atproto_client.exceptions.__name__, atproto_client.request.__name__)' + output = subprocess.run( # noqa: S603 + [sys.executable, '-c', code], capture_output=True, text=True, check=True + ).stdout + + assert output.split() == ['atproto_client.exceptions', 'atproto_client.request'] + + +def test_client_package_raises_attribute_error_for_unknown_names() -> None: + import atproto_client + + with pytest.raises(AttributeError): + _ = atproto_client.does_not_exist + + +def _always_valid(*_: object) -> bool: + return True + + +def _patched_did_key(*_: object) -> str: + return 'did:key:patched' + + +def test_crypto_aliases_stay_patchable(monkeypatch: pytest.MonkeyPatch) -> None: + from atproto_core.did_doc import did_doc + from atproto_server.auth import jwt + + monkeypatch.setattr(jwt, 'verify_signature', _always_valid) + monkeypatch.setattr(did_doc, 'get_did_key', _patched_did_key) + + assert jwt._verify_signature('key', b'input', b'signature') is True + assert did_doc.get_did_key('Multikey', 'z') == 'did:key:patched' diff --git a/tests/test_atproto/test_models.py b/tests/test_atproto/test_models.py index b1eb6401a..0b9dc2adc 100644 --- a/tests/test_atproto/test_models.py +++ b/tests/test_atproto/test_models.py @@ -14,3 +14,8 @@ def test_models_submodule_is_the_models_package() -> None: assert atproto.models is atproto_client.models assert models is atproto_client.models + + +def test_missing_dunder_raises_attribute_error() -> None: + for name in ('__qualname__', '__wrapped__', '__does_not_exist__'): + assert not hasattr(atproto_client.models, name) diff --git a/tests/test_atproto_client/models/tests/test_dot_dict.py b/tests/test_atproto_client/models/tests/test_dot_dict.py index b01a63e08..bafa03ae1 100644 --- a/tests/test_atproto_client/models/tests/test_dot_dict.py +++ b/tests/test_atproto_client/models/tests/test_dot_dict.py @@ -160,3 +160,42 @@ def test_dot_dict_mapping_protocol() -> None: assert 'a' in model assert 'b' in model assert 'c' not in model + + +def test_dot_dict_copies_nested_containers() -> None: + nested = DotDict({'x': {'y': 1}}) + leaf = bytearray(b'raw') + test_data: t.Dict[str, t.Any] = {'nested': nested, 'pair': ({'a': 1}, [2]), 'leaf': leaf} + model = DotDict(test_data) + + assert model.nested is not nested + assert model.nested.x.y == 1 + assert isinstance(model.pair, tuple) + assert model.pair[0]['a'] == 1 + assert model.leaf == leaf + assert model.leaf is not leaf + + model.nested.x.y = 2 + assert nested.x.y == 1 + + output = model.to_dict() + assert output == {'nested': {'x': {'y': 2}}, 'pair': ({'a': 1}, [2]), 'leaf': bytearray(b'raw')} + assert type(output['nested']) is dict + assert type(output['nested']['x']) is dict + assert type(output['pair'][0]) is dict + assert output['leaf'] is not model.leaf + + +def test_dot_dict_snake_case_check_matches_the_conversion() -> None: + from atproto_client.models.dot_dict import _convert_camel_case_to_snake_case, _is_snake_case + + for string in ('a', 'snake_case', 'camelCase', 'HTTPServer', 'getHTTP2', 'x1', '_private', 'émoji', 'Émoji', ''): + assert _is_snake_case(string) == (string == _convert_camel_case_to_snake_case(string)), string + + +def test_dot_dict_supports_weak_references() -> None: + import weakref + + model = DotDict({'a': 1}) + + assert weakref.ref(model)() is model diff --git a/tests/test_atproto_client/models/tests/test_shared_fields_set.py b/tests/test_atproto_client/models/tests/test_shared_fields_set.py new file mode 100644 index 000000000..bf261e984 --- /dev/null +++ b/tests/test_atproto_client/models/tests/test_shared_fields_set.py @@ -0,0 +1,87 @@ +import copy +import pickle + +import pytest +from atproto_client import models +from atproto_client.models.base import _SharedFieldsSet + +Post = models.AppBskyFeedPost.Record + + +def _post(text: str) -> models.AppBskyFeedPost.Record: + return Post.model_validate({'text': text, 'createdAt': '2024-01-01T00:00:00Z'}) + + +def test_models_of_the_same_shape_share_the_fields_set() -> None: + first, second = _post('a'), _post('b') + + assert type(first.__pydantic_fields_set__) is _SharedFieldsSet + assert first.__pydantic_fields_set__ is second.__pydantic_fields_set__ + assert second.model_dump(exclude_unset=True, by_alias=True) == {'text': 'b', 'createdAt': '2024-01-01T00:00:00Z'} + + +def test_setattr_does_not_leak_into_other_models() -> None: + first, second = _post('a'), _post('b') + + first.langs = ['en'] + first.via = 'client' + + assert 'langs' in first.__pydantic_fields_set__ + assert 'langs' not in second.__pydantic_fields_set__ + assert 'via' not in second.__pydantic_fields_set__ + assert second.model_extra == {} + + +def test_model_fields_set_is_safe_to_mutate() -> None: + first, second = _post('a'), _post('b') + + first.model_fields_set.add('langs') + + assert 'langs' not in second.__pydantic_fields_set__ + + +def test_copies_do_not_leak_into_other_models() -> None: + first, second = _post('a'), _post('b') + + updated = first.model_copy(update={'langs': ['en']}) + deep = copy.deepcopy(first) + deep.langs = ['de'] + + assert 'langs' in updated.model_fields_set + assert 'langs' not in second.__pydantic_fields_set__ + + +def test_models_pickled_together_do_not_share_a_plain_set() -> None: + first, second = pickle.loads(pickle.dumps([_post('a'), _post('b')])) # noqa: S301 + + first.langs = ['en'] + + assert 'langs' not in second.__pydantic_fields_set__ + assert 'langs' not in _post('c').__pydantic_fields_set__ + + +def test_models_with_extra_fields_keep_their_own_set() -> None: + first = Post.model_validate({'text': 'a', 'createdAt': '2024-01-01T00:00:00Z', 'via': 'client'}) + second = Post.model_validate({'text': 'a', 'createdAt': '2024-01-01T00:00:00Z', 'via': 'client'}) + + assert type(first.__pydantic_fields_set__) is set + assert first.__pydantic_fields_set__ is not second.__pydantic_fields_set__ + + +def test_deprecated_fields_set_is_safe_to_mutate() -> None: + import warnings + + first, second = _post('a'), _post('b') + + with warnings.catch_warnings(): + warnings.simplefilter('ignore') + first.__fields_set__.add('langs') + + assert 'langs' not in second.__pydantic_fields_set__ + + +def test_lexicon_warning_is_read_at_validation_time(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv('ATPROTO_LEXICON_WARN', '1') + + with pytest.warns(UserWarning, match='Extra fields found'): + Post.model_validate({'text': 'a', 'createdAt': '2024-01-01T00:00:00Z', 'via': 'client'}) diff --git a/tests/test_atproto_core/test_cid.py b/tests/test_atproto_core/test_cid.py index c001414ac..0b20aed2b 100644 --- a/tests/test_atproto_core/test_cid.py +++ b/tests/test_atproto_core/test_cid.py @@ -1,4 +1,5 @@ import libipld +import pytest from atproto_core.car import CAR from atproto_core.cid import CID, CIDType from pydantic import BaseModel @@ -110,3 +111,22 @@ def test_car_blocks_lookup_by_cid_and_str() -> None: for block_cid in car.blocks: assert car.blocks[block_cid] in (root_block, child_block) assert car.blocks[str(block_cid)] == car.blocks[block_cid] # pyrefly: ignore[bad-index] + + +def test_cid_decode_defers_the_multihash() -> None: + test_cid = 'bafyreihiurhwu7dp63wcu5x266pa7fiqm2zea7za4zfxz7hb3jhien7y5a' + decoded = libipld.decode_cid(test_cid) + + cid = CID.decode(test_cid) + assert cid._hash is None + + assert cid.version == decoded['version'] + assert cid.codec == decoded['codec'] + assert cid.hash.code == decoded['hash']['code'] + assert cid.hash.size == decoded['hash']['size'] + assert cid.hash.digest == decoded['hash']['digest'] + + +def test_cid_decode_still_validates() -> None: + with pytest.raises(ValueError): + CID.decode('not-a-cid')