From 5531dc05412dd212383a0dea01905270e10a40bd Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Fri, 24 Jul 2026 11:04:48 -0400 Subject: [PATCH 1/8] Added type hints to _converter.py --- asdf/extension/_converter.py | 59 +++++++++++++++++++++++++----------- 1 file changed, 41 insertions(+), 18 deletions(-) diff --git a/asdf/extension/_converter.py b/asdf/extension/_converter.py index 947c89ac9..fff0e167c 100644 --- a/asdf/extension/_converter.py +++ b/asdf/extension/_converter.py @@ -3,12 +3,24 @@ types. Will eventually replace the `asdf.types` module. """ +from __future__ import annotations + import abc +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable from asdf.util import get_class_name, uri_match +if TYPE_CHECKING: + from collections.abc import Iterable + + from asdf.extension import ExtensionProxy, SerializationContext + from asdf.typing import TreeKey + + YamlNode = dict[TreeKey, Any] | list[Any] | str + -class Converter(abc.ABC): +@runtime_checkable +class Converter(Protocol): """ Abstract base class for plugins that convert nodes from the parsed YAML tree into custom objects, and vice versa. @@ -44,6 +56,13 @@ class Converter(abc.ABC): information about that object to display during ``AsdfFile.info``. """ + # This is a hacky workaround for a limitation of Python protocols. + # + # @runtime_checkable protocols always support issubclass() but only support isinstance() + # if all of the class members are methods (i.e. no attributes or properties) + # + # Converter being a protocol makes it work with type checkers and issubclass(). + # This __subclasshook__ makes it also work with isinstance(). @classmethod def __subclasshook__(cls, class_): if cls is Converter: @@ -57,7 +76,7 @@ def __subclasshook__(cls, class_): @property @abc.abstractmethod - def tags(self): + def tags(self) -> Iterable[str]: """ Get the YAML tags that this converter is capable of handling. URI patterns are permitted, see @@ -68,10 +87,11 @@ def tags(self): iterable of str Tag URIs or URI patterns. """ + ... @property @abc.abstractmethod - def types(self): + def types(self) -> Iterable[str | type]: """ Get the Python types that this converter is capable of handling. @@ -81,9 +101,10 @@ def types(self): iterable of str or type If str, the fully qualified class name of the type. """ + ... @abc.abstractmethod - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNode: """ Convert an object into a node suitable for YAML serialization. This method is not responsible for writing actual YAML; rather, it @@ -117,9 +138,10 @@ def to_yaml_tree(self, obj, tag, ctx): dict or list or str The YAML node representation of the object. """ + ... @abc.abstractmethod - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> Any: """ Convert a YAML node into an instance of a custom type. @@ -152,6 +174,7 @@ def from_yaml_tree(self, node, tag, ctx): An instance of one of the types listed in the `types` property, or a generator that yields such an instance. """ + ... class ConverterProxy(Converter): @@ -160,7 +183,7 @@ class ConverterProxy(Converter): implementations of optional methods. """ - def __init__(self, delegate, extension): + def __init__(self, delegate: Converter, extension: ExtensionProxy): if not isinstance(delegate, Converter): msg = "Converter must implement the asdf.extension.Converter interface" raise TypeError(msg) @@ -201,7 +224,7 @@ def __init__(self, delegate, extension): raise TypeError(msg) @property - def lazy(self): + def lazy(self) -> bool: """ Boolean indicating if this Converter supports "lazy" node objects @@ -212,7 +235,7 @@ def lazy(self): return getattr(self._delegate, "lazy", False) @property - def tags(self): + def tags(self) -> list[str]: """ Get the list of tag URIs that this converter is capable of handling. @@ -224,7 +247,7 @@ def tags(self): return self._tags @property - def types(self): + def types(self) -> list[str | type]: """ Get the Python types that this converter is capable of handling. @@ -235,7 +258,7 @@ def types(self): """ return self._types - def select_tag(self, obj, ctx): + def select_tag(self, obj: Any, ctx: SerializationContext) -> str | None: """ Select the tag to use when converting an object to YAML. @@ -257,7 +280,7 @@ def select_tag(self, obj, ctx): return method(obj, self._tags, ctx) - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNode: """ Convert an object into a node suitable for YAML serialization. @@ -278,7 +301,7 @@ def to_yaml_tree(self, obj, tag, ctx): """ return self._delegate.to_yaml_tree(obj, tag, ctx) - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> Any: """ Convert a YAML node into an instance of a custom type. @@ -297,7 +320,7 @@ def from_yaml_tree(self, node, tag, ctx): """ return self._delegate.from_yaml_tree(node, tag, ctx) - def to_info(self, obj): + def to_info(self, obj: Any) -> Any: """ Convert an object to a container with items further defining information about this node. This method @@ -319,7 +342,7 @@ def to_info(self, obj): return self._delegate.to_info(obj) @property - def delegate(self): + def delegate(self) -> Converter: """ Get the wrapped converter instance. @@ -330,7 +353,7 @@ def delegate(self): return self._delegate @property - def extension(self): + def extension(self) -> ExtensionProxy: """ Get the extension that provided this converter. @@ -341,7 +364,7 @@ def extension(self): return self._extension @property - def package_name(self): + def package_name(self) -> str | None: """ Get the name of the Python package of this converter's extension. This may not be the same package that implements @@ -355,7 +378,7 @@ def package_name(self): return self.extension.package_name @property - def package_version(self): + def package_version(self) -> str | None: """ Get the version of the Python package of this converter's extension. This may not be the same package that implements @@ -369,7 +392,7 @@ def package_version(self): return self.extension.package_version @property - def class_name(self): + def class_name(self) -> str: """ Get the fully qualified class name of this converter. From dbc6453b95ecc3b36f0b1a544401863cd2acf5bb Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Fri, 31 Jul 2026 13:26:29 -0400 Subject: [PATCH 2/8] Added type hints --- asdf/_asdf.py | 3 +- asdf/_block/manager.py | 6 +- asdf/_block/writer.py | 6 +- asdf/_compression.py | 75 ++++++++++++----- asdf/_core/_converters/complex.py | 12 ++- asdf/_core/_converters/constant.py | 17 ++-- asdf/_core/_converters/external_reference.py | 23 ++++-- asdf/_core/_converters/integer.py | 28 ++++--- asdf/_core/_converters/ndarray.py | 41 +++++++--- asdf/_core/_converters/reference.py | 18 +++- asdf/_core/_converters/tree.py | 44 +++++----- asdf/_helpers.py | 9 +- asdf/_node_info.py | 17 ++-- asdf/_tests/test_extension.py | 8 +- asdf/config.py | 3 +- asdf/extension/__init__.py | 8 +- asdf/extension/_compressor.py | 56 +++++++++---- asdf/extension/_converter.py | 47 ++++------- asdf/extension/_extension.py | 86 ++++++++++++-------- asdf/extension/_manager.py | 34 ++++---- asdf/extension/_manifest.py | 38 ++++++--- asdf/extension/_serialization_context.py | 64 +++++++++------ asdf/extension/_tag.py | 19 +++-- asdf/extension/_validator.py | 6 +- asdf/tags/core/integer.py | 20 ++++- asdf/tags/core/ndarray.py | 54 ++++++------ asdf/tags/core/stream.py | 2 +- asdf/typing.py | 29 +++---- asdf/versioning.py | 12 ++- 29 files changed, 493 insertions(+), 292 deletions(-) diff --git a/asdf/_asdf.py b/asdf/_asdf.py index 9801bd693..1297d8f07 100644 --- a/asdf/_asdf.py +++ b/asdf/_asdf.py @@ -35,14 +35,13 @@ from collections.abc import Mapping, MutableMapping, Sequence from typing import Any - from asdf.extension import ExtensionManager, SerializationContext + from asdf.extension import ExtensionLike, ExtensionManager, SerializationContext from asdf.generic_io import GenericFile from asdf.tagged import Tagged from asdf.typing import ( ArrayStorage, AsdfVersionLike, Compression, - ExtensionLike, FileLike, FileMode, FilterFn, diff --git a/asdf/_block/manager.py b/asdf/_block/manager.py index 97a37f4fb..0e4211028 100644 --- a/asdf/_block/manager.py +++ b/asdf/_block/manager.py @@ -20,7 +20,7 @@ from asdf._block.key import Key from asdf.generic_io import GenericFile - from asdf.typing import ArrayStorage, BlockDataCallback, ByteArray1D, Compression, NDArray + from asdf.typing import ArrayCallback, ArrayStorage, ByteArray1D, Compression, NDArray class ReadBlocks(collections.UserList[ReadBlock]): @@ -402,7 +402,7 @@ def _write_external_blocks(self, write_checksums: bool) -> None: af.write_to(f, include_block_index=False) writer.write_blocks(f, [blk], write_checksums=write_checksums) - def make_write_block(self, data: ByteArray1D | BlockDataCallback, options: Options | None, obj: Any) -> int | str: + def make_write_block(self, data: NDArray | ArrayCallback, options: Options | None, obj: Any) -> int | str: """ Make a WriteBlock with data and options and associate it with an object (obj). @@ -465,7 +465,7 @@ def make_write_block(self, data: ByteArray1D | BlockDataCallback, options: Optio index = self._write_blocks.append_block(blk, obj) return index - def set_streamed_write_block(self, data: ByteArray1D | BlockDataCallback, obj: Any) -> None: + def set_streamed_write_block(self, data: NDArray | ArrayCallback | None, obj: Any) -> None: """ Create a WriteBlock that will be written as an ASDF streamed block. diff --git a/asdf/_block/writer.py b/asdf/_block/writer.py index 61fe129ac..d7fbf5a89 100644 --- a/asdf/_block/writer.py +++ b/asdf/_block/writer.py @@ -13,7 +13,7 @@ from asdf._block.io import BlockHeader from asdf.generic_io import GenericFile - from asdf.typing import BlockDataCallback, ByteArray1D, Compression + from asdf.typing import ArrayCallback, ByteArray1D, Compression, NDArray class WriteBlock: @@ -25,7 +25,7 @@ class WriteBlock: def __init__( self, - data: ByteArray1D | BlockDataCallback | None, + data: NDArray | ArrayCallback | None, compression: Compression = None, compression_kwargs: dict[str, Any] | None = None, ): @@ -34,7 +34,7 @@ def __init__( self.compression_kwargs: dict[str, Any] | None = compression_kwargs @property - def data(self) -> ByteArray1D | None: + def data(self) -> NDArray | None: if callable(self._data): return self._data() return self._data diff --git a/asdf/_compression.py b/asdf/_compression.py index d1c6d8bb3..2fb7b8b31 100644 --- a/asdf/_compression.py +++ b/asdf/_compression.py @@ -5,18 +5,23 @@ import typing import warnings import zlib -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, overload import numpy as np +from asdf.extension import Compress, CompressionPlugin, Decompress + from .config import get_config from .exceptions import AsdfWarning if TYPE_CHECKING: + from collections.abc import Iterable, Iterator from io import IOBase + from asdf.extension import ExtensionProxy from asdf.generic_io import GenericFile - from asdf.typing import ByteArray1D, Compression + from asdf.typing import ByteArray1D + from asdf.typing import Compression as CompressionLabel def validate(compression: str | bytes | None) -> str | None: @@ -62,7 +67,9 @@ def validate(compression: str | bytes | None) -> str | None: return compression -class Lz4Compressor: +class Lz4Compressor(Compress, Decompress): + label = b"lz4" + def __init__(self): try: import lz4.block @@ -76,7 +83,7 @@ def __init__(self): self._api = lz4.block - def compress(self, data, **kwargs): + def compress(self, data: memoryview, **kwargs) -> Iterator[bytes]: kwargs["mode"] = kwargs.get("mode", "default") compression_block_size = kwargs.pop("compression_block_size", 1 << 22) @@ -86,14 +93,14 @@ def compress(self, data, **kwargs): header = struct.pack("!I", len(_output)) yield header + _output - def decompress(self, blocks, out, **kwargs): + def decompress(self, data: Iterable[bytes], out, **kwargs) -> int: _size = 0 _pos = 0 _partial_len = b"" _buffer = None bytesout = 0 - for block in blocks: + for block in data: cast = "c" blk = memoryview(block).cast(cast) # don't copy on slice @@ -149,39 +156,57 @@ def decompress(self, blocks, out, **kwargs): return bytesout -class ZlibCompressor: - def compress(self, data, **kwargs): +class ZlibCompressor(Compress, Decompress): + label = b"zlib" + + def compress(self, data: memoryview, **kwargs) -> Iterator[bytes]: comp = zlib.compress(data, **kwargs) yield comp - def decompress(self, blocks, out, **kwargs): + def decompress(self, data: Iterable[bytes], out, **kwargs) -> int: decompressor = zlib.decompressobj(**kwargs) i = 0 - for block in blocks: + for block in data: decomp = decompressor.decompress(block) out[i : i + len(decomp)] = decomp i += len(decomp) return i -class Bzp2Compressor: - def compress(self, data, **kwargs): +class Bzp2Compressor(Compress, Decompress): + label = b"bzp2" + + def compress(self, data: memoryview, **kwargs) -> Iterator[bytes]: comp = bz2.compress(data, **kwargs) yield comp - def decompress(self, blocks, out, **kwargs): + def decompress(self, data: Iterable[bytes], out, **kwargs) -> int: decompressor = bz2.BZ2Decompressor(**kwargs) i = 0 - for block in blocks: + for block in data: decomp = decompressor.decompress(block) out[i : i + len(decomp)] = decomp i += len(decomp) return i -def _get_compressor_from_extensions(compression, return_extension=False): +@overload +def _get_compressor_from_extensions( + compression: bytes | str | None, return_extension: Literal[True] +) -> tuple[CompressionPlugin, ExtensionProxy] | None: ... +@overload +def _get_compressor_from_extensions( + compression: bytes | str | None, return_extension: Literal[False] +) -> CompressionPlugin | None: ... +@overload +def _get_compressor_from_extensions(compression: bytes | str | None) -> CompressionPlugin | None: ... + + +def _get_compressor_from_extensions( + compression: bytes | str | None, return_extension: bool = False +) -> CompressionPlugin | tuple[CompressionPlugin, ExtensionProxy] | None: """ Look at the loaded ASDF extensions and return the first one (if any) that can handle this type of compression. @@ -202,7 +227,7 @@ def _get_compressor_from_extensions(compression, return_extension=False): return None -def _get_all_compression_extension_labels(): +def _get_all_compression_extension_labels() -> list[str]: """ Get the list of compression labels supported via extensions """ @@ -215,12 +240,12 @@ def _get_all_compression_extension_labels(): return labels -def _get_compressor(label: str) -> Any: +def _get_compressor(label: str) -> CompressionPlugin: ext_comp = _get_compressor_from_extensions(label) if ext_comp is not None: # Use an extension before builtins - comp = ext_comp + return ext_comp elif label == "zlib": comp = ZlibCompressor() elif label == "bzp2": @@ -234,7 +259,7 @@ def _get_compressor(label: str) -> Any: return comp -def to_compression_header(compression: Compression) -> bytes: +def to_compression_header(compression: CompressionLabel) -> bytes: """ Converts a compression string to the four byte field in a block header. @@ -281,6 +306,10 @@ def decompress( compression = typing.cast("str", validate(compression)) decoder = _get_compressor(compression) + if not isinstance(decoder, Decompress): + msg = f"Compression plugin {decoder.label} does not implement decompress() function" + raise TypeError(msg) + if config is None: config = {} @@ -320,6 +349,10 @@ def compress( """ compression = typing.cast("str", validate(compression)) encoder = _get_compressor(compression) + if not isinstance(encoder, Compress): + msg = f"Compression plugin {encoder.label} does not implement compress() function" + raise TypeError(msg) + if config is None: config = {} @@ -335,12 +368,12 @@ def compress( # get a 1D array that preserves byteorder # Note: in Python < 3.12 numpy typing doesn't correctly reflect that arrays support buffer protocol # See also: https://github.com/numpy/numpy/issues/26783 - view = memoryview(np.frombuffer(view, dtype=view.format)) # pyrefly: ignore[bad-argument-type] + view: memoryview = memoryview(np.frombuffer(view, dtype=view.format)) # pyrefly: ignore[bad-argument-type] if not view.contiguous: # the data will be contiguous by construction, but better safe than sorry! raise ValueError(view.contiguous) - compressed = encoder.compress(data, **config) + compressed = encoder.compress(view, **config) # Write block by block for comp in compressed: fd.write(comp) diff --git a/asdf/_core/_converters/complex.py b/asdf/_core/_converters/complex.py index b62e7d58c..a9e340e98 100644 --- a/asdf/_core/_converters/complex.py +++ b/asdf/_core/_converters/complex.py @@ -1,24 +1,30 @@ +from __future__ import annotations + import re +from typing import TYPE_CHECKING import numpy as np from asdf import util from asdf.extension import Converter +if TYPE_CHECKING: + from asdf.extension import SerializationContext + _REPLACEMENTS = { re.compile("i(?!nf)"): "j", re.compile("I(?!NF)"): "J", } -class ComplexConverter(Converter): +class ComplexConverter(Converter[complex, str]): tags = ["tag:stsci.edu:asdf/core/complex-1.0.0"] types = [*list(util._iter_subclasses(np.complexfloating)), complex] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: complex, tag: str, ctx: SerializationContext) -> str: return str(obj) - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: str, tag: str, ctx: SerializationContext) -> complex: for pattern, replacement in _REPLACEMENTS.items(): node = pattern.sub(replacement, node) diff --git a/asdf/_core/_converters/constant.py b/asdf/_core/_converters/constant.py index bcec02ec3..386c99e60 100644 --- a/asdf/_core/_converters/constant.py +++ b/asdf/_core/_converters/constant.py @@ -1,14 +1,21 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + from asdf.extension import Converter +from asdf.tags.core import Constant +if TYPE_CHECKING: + from asdf.extension import SerializationContext + from asdf.typing import YamlNode -class ConstantConverter(Converter): + +class ConstantConverter(Converter[Constant]): tags = ["tag:stsci.edu:asdf/core/constant-1.0.0"] types = ["asdf.tags.core.constant.Constant"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: Constant, tag: str, ctx: SerializationContext) -> YamlNode: return obj.value - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import Constant - + def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> Constant: return Constant(node) diff --git a/asdf/_core/_converters/external_reference.py b/asdf/_core/_converters/external_reference.py index 81198b7b8..d583639a7 100644 --- a/asdf/_core/_converters/external_reference.py +++ b/asdf/_core/_converters/external_reference.py @@ -1,11 +1,26 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypedDict + from asdf.extension import Converter +from asdf.tags.core import ExternalArrayReference + +if TYPE_CHECKING: + from asdf.extension import SerializationContext -class ExternalArrayReferenceConverter(Converter): +class ExternalArrayNode(TypedDict): + fileuri: str + target: Any + datatype: str + shape: list[str] + + +class ExternalArrayReferenceConverter(Converter[ExternalArrayReference, ExternalArrayNode]): tags = ["tag:stsci.edu:asdf/core/externalarray-1.0.0"] types = ["asdf.tags.core.external_reference.ExternalArrayReference"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: ExternalArrayReference, tag: str, ctx: SerializationContext) -> ExternalArrayNode: return { "fileuri": obj.fileuri, "target": obj.target, @@ -13,7 +28,5 @@ def to_yaml_tree(self, obj, tag, ctx): "shape": list(obj.shape), } - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import ExternalArrayReference - + def from_yaml_tree(self, node: ExternalArrayNode, tag: str, ctx: SerializationContext) -> ExternalArrayReference: return ExternalArrayReference(node["fileuri"], node["target"], node["datatype"], tuple(node["shape"])) diff --git a/asdf/_core/_converters/integer.py b/asdf/_core/_converters/integer.py index cce8b36b0..9a74b1a51 100644 --- a/asdf/_core/_converters/integer.py +++ b/asdf/_core/_converters/integer.py @@ -1,9 +1,23 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypedDict + import numpy as np from asdf.extension import Converter +from asdf.tags.core.integer import IntegerType + +if TYPE_CHECKING: + from asdf.extension import SerializationContext -class IntegerConverter(Converter): +class IntegerNode(TypedDict): + words: Any + sign: str + string: str + + +class IntegerConverter(Converter[IntegerType, IntegerNode]): tags = [ "tag:stsci.edu:asdf/core/integer-1.0.0", "tag:stsci.edu:asdf/core/integer-1.1.0", @@ -11,7 +25,7 @@ class IntegerConverter(Converter): ] types = ["asdf.tags.core.integer.IntegerType"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: IntegerType, tag: str, ctx: SerializationContext) -> IntegerNode: abs_value = int(np.abs(obj._value)) # pack integer value into 32-bit words @@ -23,17 +37,11 @@ def to_yaml_tree(self, obj, tag, ctx): array = np.array(words, dtype=np.uint32) - tree = {} ctx.set_array_storage(array, obj._storage) - tree["words"] = array - tree["sign"] = obj._sign - tree["string"] = str(int(obj._value)) - - return tree - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core.integer import IntegerType + return {"words": array, "sign": obj._sign, "string": str(int(obj._value))} + def from_yaml_tree(self, node: IntegerNode, tag: str, ctx: SerializationContext) -> IntegerType: value = 0 for x in node["words"][::-1]: value <<= 32 diff --git a/asdf/_core/_converters/ndarray.py b/asdf/_core/_converters/ndarray.py index 8b301ba3d..655dc0756 100644 --- a/asdf/_core/_converters/ndarray.py +++ b/asdf/_core/_converters/ndarray.py @@ -1,9 +1,30 @@ +from typing import Any, Literal + import numpy as np +from typing_extensions import TypedDict + +from asdf.extension import Converter, SerializationContext +from asdf.tags.core.ndarray import NDArrayType +from asdf.tags.core.stream import Stream +from asdf.typing import NDArray + -from asdf.extension import Converter +class NdArrayMap(TypedDict, total=False): + source: str | int + data: Any + datatype: str + shape: list[str | int] + byteorder: Literal["little", "big"] + offset: int + strides: list[int] + mask: Any -class NDArrayConverter(Converter): +_NdArray = NDArray | Stream +_NdArrayNode = NdArrayMap | list[Any] + + +class NDArrayConverter(Converter[_NdArray, _NdArrayNode]): tags = [ "tag:stsci.edu:asdf/core/ndarray-1.0.0", "tag:stsci.edu:asdf/core/ndarray-1.1.0", @@ -17,22 +38,20 @@ class NDArrayConverter(Converter): "asdf.tags.core.stream.Stream", ] - def to_yaml_tree(self, obj, tag, ctx): - import numpy as np + def to_yaml_tree(self, obj: _NdArray, tag: str, ctx: SerializationContext) -> _NdArrayNode: from numpy import ma from asdf import config, util from asdf._block.options import Options - from asdf.tags.core.ndarray import NDArrayType, numpy_array_to_list, numpy_dtype_to_asdf_datatype - from asdf.tags.core.stream import Stream + from asdf.tags.core.ndarray import numpy_array_to_list, numpy_dtype_to_asdf_datatype data = obj + result: NdArrayMap = {} - if isinstance(obj, Stream): + if isinstance(data, Stream): # previously, stream never passed on data, we can do that here ctx._blocks.set_streamed_write_block(data._array, data) - result = {} result["source"] = -1 result["shape"] = ["*", *data._shape] result["datatype"] = data._datatype @@ -98,8 +117,6 @@ def to_yaml_tree(self, obj, tag, ctx): include_byteorder=(options.storage_type != "inline"), ) - result = {} - result["shape"] = list(shape) if options.storage_type == "streamed": result["shape"][0] = "*" @@ -134,7 +151,7 @@ def to_yaml_tree(self, obj, tag, ctx): return result - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: _NdArrayNode, tag: str, ctx: SerializationContext) -> _NdArray: import sys import weakref @@ -194,5 +211,5 @@ def data_callback(_attr=None, _ref=weakref.ref(ctx._blocks)): msg = "Invalid ndarray description." raise TypeError(msg) - def to_info(self, obj): + def to_info(self, obj: NDArray) -> dict[str, Any]: return {"shape": obj.shape, "dtype": obj.dtype} diff --git a/asdf/_core/_converters/reference.py b/asdf/_core/_converters/reference.py index eeedaec25..8bfd7d39a 100644 --- a/asdf/_core/_converters/reference.py +++ b/asdf/_core/_converters/reference.py @@ -1,11 +1,21 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, TypedDict + from asdf.extension import Converter +from asdf.reference import Reference + +if TYPE_CHECKING: + from asdf.extension import SerializationContext + +ReferenceNode = TypedDict("ReferenceNode", {"$ref": str}) -class ReferenceConverter(Converter): +class ReferenceConverter(Converter[Reference, ReferenceNode]): tags = [] types = ["asdf.reference.Reference"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: Reference, tag: str, ctx: SerializationContext) -> ReferenceNode: from asdf.generic_io import relative_uri base_uri = None @@ -16,8 +26,8 @@ def to_yaml_tree(self, obj, tag, ctx): uri = relative_uri(base_uri, obj._uri) if base_uri is not None else obj._uri return {"$ref": uri} - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: ReferenceNode, tag: str, ctx: SerializationContext) -> Reference: raise NotImplementedError() - def select_tag(self, obj, tags, ctx): + def select_tag(self, obj: Reference, tags: list[str], ctx: SerializationContext) -> str | None: return None diff --git a/asdf/_core/_converters/tree.py b/asdf/_core/_converters/tree.py index 679f09fb3..faa6b16e8 100644 --- a/asdf/_core/_converters/tree.py +++ b/asdf/_core/_converters/tree.py @@ -1,7 +1,13 @@ -from asdf.extension import Converter +from typing import Any +from asdf.extension import Converter, SerializationContext +from asdf.tags.core import AsdfObject, ExtensionMetadata, HistoryEntry, Software, SubclassMetadata +from asdf.typing import TreeKey -class AsdfObjectConverter(Converter): +_YamlMap = dict[TreeKey, Any] + + +class AsdfObjectConverter(Converter[AsdfObject, _YamlMap]): # Since AsdfObject is just a dict, we're able to use the same converter # for both tag versions. tags = [ @@ -10,42 +16,36 @@ class AsdfObjectConverter(Converter): ] types = ["asdf.tags.core.AsdfObject"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: AsdfObject, tag: str, ctx: SerializationContext) -> _YamlMap: return dict(obj) - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import AsdfObject - + def from_yaml_tree(self, node: _YamlMap, tag: str, ctx: SerializationContext) -> AsdfObject: return AsdfObject(node) -class ExtensionMetadataConverter(Converter): +class ExtensionMetadataConverter(Converter[ExtensionMetadata, _YamlMap]): tags = ["tag:stsci.edu:asdf/core/extension_metadata-1.0.0"] types = ["asdf.tags.core.ExtensionMetadata"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: ExtensionMetadata, tag: str, ctx: SerializationContext) -> _YamlMap: return dict(obj) - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import ExtensionMetadata - + def from_yaml_tree(self, node: _YamlMap, tag: str, ctx: SerializationContext) -> ExtensionMetadata: return ExtensionMetadata(node) -class HistoryEntryConverter(Converter): +class HistoryEntryConverter(Converter[HistoryEntry, _YamlMap]): tags = ["tag:stsci.edu:asdf/core/history_entry-1.0.0"] types = ["asdf.tags.core.HistoryEntry"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: HistoryEntry, tag: str, ctx: SerializationContext) -> _YamlMap: return dict(obj) - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import HistoryEntry - + def from_yaml_tree(self, node: _YamlMap, tag: str, ctx: SerializationContext) -> HistoryEntry: return HistoryEntry(node) -class SoftwareConverter(Converter): +class SoftwareConverter(Converter[Software, _YamlMap]): tags = ["tag:stsci.edu:asdf/core/software-1.0.0"] types = ["asdf.tags.core.Software"] @@ -53,19 +53,15 @@ def to_yaml_tree(self, obj, tag, ctx): return dict(obj) def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import Software - return Software(node) -class SubclassMetadataConverter(Converter): +class SubclassMetadataConverter(Converter[SubclassMetadata, _YamlMap]): tags = ["tag:stsci.edu:asdf/core/subclass_metadata-1.0.0"] types = ["asdf.tags.core.SubclassMetadata"] - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: SubclassMetadata, tag: str, ctx: SerializationContext) -> _YamlMap: return dict(obj) - def from_yaml_tree(self, node, tag, ctx): - from asdf.tags.core import SubclassMetadata - + def from_yaml_tree(self, node: _YamlMap, tag: str, ctx: SerializationContext) -> SubclassMetadata: return SubclassMetadata(node) diff --git a/asdf/_helpers.py b/asdf/_helpers.py index 40b4be013..cbef50005 100644 --- a/asdf/_helpers.py +++ b/asdf/_helpers.py @@ -1,8 +1,15 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + from . import versioning from ._version import version as asdf_package_version +if TYPE_CHECKING: + from asdf.versioning import AsdfVersion + -def validate_version(version): +def validate_version(version: str | AsdfVersion) -> str: # Account for the possibility of AsdfVersion version = str(version) if version not in versioning.supported_versions: diff --git a/asdf/_node_info.py b/asdf/_node_info.py index 604e72965..87634f131 100644 --- a/asdf/_node_info.py +++ b/asdf/_node_info.py @@ -1,9 +1,16 @@ +from __future__ import annotations + import re -from collections import namedtuple +from typing import TYPE_CHECKING, Any + +from typing_extensions import NamedTuple from .schema import load_schema from .treeutil import get_children, is_container +if TYPE_CHECKING: + from asdf.typing import TreeKey + def _filter_tree(info, filters): """ @@ -231,10 +238,7 @@ def _make_traversable(node, extension_manager): return extension_manager.get_converter_for_type(node_type).to_info(node), False, True -_SchemaInfo = namedtuple("SchemaInfo", ["info", "value"]) - - -class SchemaInfo(_SchemaInfo): +class SchemaInfo(NamedTuple): """ A class to hold the schema info and the value of the node. @@ -246,6 +250,9 @@ class SchemaInfo(_SchemaInfo): The value of the node. """ + info: dict[TreeKey, Any] + value: Any + def __repr__(self): return f"{self.info}" diff --git a/asdf/_tests/test_extension.py b/asdf/_tests/test_extension.py index ae13a46ad..28f835aa8 100644 --- a/asdf/_tests/test_extension.py +++ b/asdf/_tests/test_extension.py @@ -452,13 +452,15 @@ class ConverterNoSubclass: tags = [] types = [] - def to_yaml_tree(self, *args): + def to_yaml_tree(self, obj, tag, ctx): pass - def from_yaml_tree(self, *args): + def from_yaml_tree(self, node, tag, ctx): pass - assert issubclass(ConverterNoSubclass, Converter) + # Have to use isinstance instead of issubclass + # issubclass isn't supported for Protocols with non-function attributes + assert isinstance(ConverterNoSubclass(), Converter) def test_converter_proxy(): diff --git a/asdf/config.py b/asdf/config.py index 59cc3e756..b27fbbf87 100644 --- a/asdf/config.py +++ b/asdf/config.py @@ -19,7 +19,8 @@ if TYPE_CHECKING: from collections.abc import Generator, Mapping - from asdf.typing import ArrayStorage, Compression, ExtensionLike + from asdf.extension import ExtensionLike + from asdf.typing import ArrayStorage, Compression __all__ = ["AsdfConfig", "config_context", "get_config"] diff --git a/asdf/extension/__init__.py b/asdf/extension/__init__.py index c633f658f..590f60642 100644 --- a/asdf/extension/__init__.py +++ b/asdf/extension/__init__.py @@ -3,9 +3,9 @@ additional custom types. """ -from ._compressor import Compressor +from ._compressor import Compress, CompressionPlugin, Compressor, Decompress from ._converter import Converter, ConverterProxy -from ._extension import Extension, ExtensionProxy +from ._extension import Extension, ExtensionLike, ExtensionProxy from ._manager import ExtensionManager, get_cached_extension_manager from ._manifest import ManifestExtension from ._serialization_context import SerializationContext @@ -14,10 +14,14 @@ __all__ = [ # New API + "Compress", + "CompressionPlugin", "Compressor", "Converter", "ConverterProxy", + "Decompress", "Extension", + "ExtensionLike", "ExtensionManager", "ExtensionProxy", "ManifestExtension", diff --git a/asdf/extension/_compressor.py b/asdf/extension/_compressor.py index cc9d4f8e4..a5cb8050a 100644 --- a/asdf/extension/_compressor.py +++ b/asdf/extension/_compressor.py @@ -9,27 +9,26 @@ of custom Python types into the YAML tree. """ +from __future__ import annotations + import abc +from typing import TYPE_CHECKING, Protocol, runtime_checkable +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator -class Compressor(abc.ABC): - """ - Abstract base class for plugins that compress binary data. - Implementing classes must provide the ``labels`` property, and - at least one of the `compress()` and `decompress()` methods. - May also provide a constructor. - """ +@runtime_checkable +class CompressionPlugin(Protocol): + """A compression plugin with an associated label. - @classmethod - def __subclasshook__(cls, class_): - if cls is Compressor: - return hasattr(class_, "label") and (hasattr(class_, "compress") or hasattr(class_, "decompress")) - return NotImplemented # pragma: no cover + For implementations of the actual compress/decompress methods, + see the `Compress` and `Decompress` protocols. + """ @property @abc.abstractmethod - def label(self): + def label(self) -> bytes: """ Get the 4-byte label identifying this compression @@ -38,8 +37,14 @@ def label(self): label : bytes The compression label """ + ... + - def compress(self, data, **kwargs): +@runtime_checkable +class Compress(CompressionPlugin, Protocol): + """A compression plugin that implements ``compress``.""" + + def compress(self, data: memoryview, **kwargs) -> Iterator[bytes]: """ Compress ``data``, yielding the results. The yield may be block-by-block, or all at once. @@ -60,7 +65,12 @@ def compress(self, data, **kwargs): """ raise NotImplementedError - def decompress(self, data, out, **kwargs): + +@runtime_checkable +class Decompress(CompressionPlugin, Protocol): + """A compression plugin that implements ``decompress``.""" + + def decompress(self, data: Iterable[bytes], out: memoryview, **kwargs) -> int: """ Decompress ``data``, writing the result into ``out``. @@ -82,3 +92,19 @@ def decompress(self, data, out, **kwargs): The number of bytes written to ``out`` """ raise NotImplementedError + + +class Compressor(Compress, Decompress): + """ + Abstract base class for plugins that compress binary data. + + Implementing classes must provide the ``labels`` property, and + at least one of the `compress()` and `decompress()` methods. + May also provide a constructor. + """ + + @classmethod + def __subclasshook__(cls, class_): + if cls is Compressor: + return hasattr(class_, "label") and (hasattr(class_, "compress") or hasattr(class_, "decompress")) + return NotImplemented # pragma: no cover diff --git a/asdf/extension/_converter.py b/asdf/extension/_converter.py index fff0e167c..a4c685a0b 100644 --- a/asdf/extension/_converter.py +++ b/asdf/extension/_converter.py @@ -6,21 +6,24 @@ from __future__ import annotations import abc -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Generic, Protocol, runtime_checkable +from typing_extensions import TypeVar + +from asdf.typing import YamlNode from asdf.util import get_class_name, uri_match if TYPE_CHECKING: from collections.abc import Iterable from asdf.extension import ExtensionProxy, SerializationContext - from asdf.typing import TreeKey - YamlNode = dict[TreeKey, Any] | list[Any] | str +_T = TypeVar("_T") +_Node = TypeVar("_Node", default=YamlNode) @runtime_checkable -class Converter(Protocol): +class Converter(Protocol[_T, _Node]): """ Abstract base class for plugins that convert nodes from the parsed YAML tree into custom objects, and vice versa. @@ -56,24 +59,6 @@ class Converter(Protocol): information about that object to display during ``AsdfFile.info``. """ - # This is a hacky workaround for a limitation of Python protocols. - # - # @runtime_checkable protocols always support issubclass() but only support isinstance() - # if all of the class members are methods (i.e. no attributes or properties) - # - # Converter being a protocol makes it work with type checkers and issubclass(). - # This __subclasshook__ makes it also work with isinstance(). - @classmethod - def __subclasshook__(cls, class_): - if cls is Converter: - return ( - hasattr(class_, "tags") - and hasattr(class_, "types") - and hasattr(class_, "to_yaml_tree") - and hasattr(class_, "from_yaml_tree") - ) - return NotImplemented # pragma: no cover - @property @abc.abstractmethod def tags(self) -> Iterable[str]: @@ -104,7 +89,7 @@ def types(self) -> Iterable[str | type]: ... @abc.abstractmethod - def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNode: + def to_yaml_tree(self, obj: _T, tag: str, ctx: SerializationContext) -> _Node: """ Convert an object into a node suitable for YAML serialization. This method is not responsible for writing actual YAML; rather, it @@ -141,7 +126,7 @@ def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNod ... @abc.abstractmethod - def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> Any: + def from_yaml_tree(self, node: _Node, tag: str, ctx: SerializationContext) -> _T: """ Convert a YAML node into an instance of a custom type. @@ -177,13 +162,13 @@ def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> ... -class ConverterProxy(Converter): +class ConverterProxy(Generic[_T, _Node], Converter[_T, _Node]): """ Proxy that wraps a `Converter` and provides default implementations of optional methods. """ - def __init__(self, delegate: Converter, extension: ExtensionProxy): + def __init__(self, delegate: Converter[_T, _Node], extension: ExtensionProxy): if not isinstance(delegate, Converter): msg = "Converter must implement the asdf.extension.Converter interface" raise TypeError(msg) @@ -258,7 +243,7 @@ def types(self) -> list[str | type]: """ return self._types - def select_tag(self, obj: Any, ctx: SerializationContext) -> str | None: + def select_tag(self, obj: _T, ctx: SerializationContext) -> str | None: """ Select the tag to use when converting an object to YAML. @@ -280,7 +265,7 @@ def select_tag(self, obj: Any, ctx: SerializationContext) -> str | None: return method(obj, self._tags, ctx) - def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNode: + def to_yaml_tree(self, obj: _T, tag: str, ctx: SerializationContext) -> _Node: """ Convert an object into a node suitable for YAML serialization. @@ -301,7 +286,7 @@ def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> YamlNod """ return self._delegate.to_yaml_tree(obj, tag, ctx) - def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> Any: + def from_yaml_tree(self, node: _Node, tag: str, ctx: SerializationContext) -> _T: """ Convert a YAML node into an instance of a custom type. @@ -320,7 +305,7 @@ def from_yaml_tree(self, node: YamlNode, tag: str, ctx: SerializationContext) -> """ return self._delegate.from_yaml_tree(node, tag, ctx) - def to_info(self, obj: Any) -> Any: + def to_info(self, obj: _T) -> Any: """ Convert an object to a container with items further defining information about this node. This method @@ -342,7 +327,7 @@ def to_info(self, obj: Any) -> Any: return self._delegate.to_info(obj) @property - def delegate(self) -> Converter: + def delegate(self) -> Converter[_T, _Node]: """ Get the wrapped converter instance. diff --git a/asdf/extension/_extension.py b/asdf/extension/_extension.py index f1ae6868b..126d5e72a 100644 --- a/asdf/extension/_extension.py +++ b/asdf/extension/_extension.py @@ -1,6 +1,7 @@ from __future__ import annotations import abc +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable from packaging.specifiers import SpecifierSet @@ -11,20 +12,17 @@ from ._tag import TagDefinition from ._validator import Validator +if TYPE_CHECKING: + from collections.abc import Iterable -class Extension(abc.ABC): - """ - Abstract base class defining an extension to ASDF. + from asdf.extension import Converter - Implementing classes must provide the `extension_uri`. - Other properties are optional. - """ - @classmethod - def __subclasshook__(cls, class_): - if cls is Extension: - return hasattr(class_, "extension_uri") - return NotImplemented # pragma: no cover +# Alternate version of `Extension` for use in type-hints +# The way `Extension` works is weird enough that it can't be replaced in actual code without a lot of changes +@runtime_checkable +class ExtensionLike(Protocol): + """Object that contains an extension URI and can be wrapped by ``ExtensionProxy``.""" @property @abc.abstractmethod @@ -38,9 +36,25 @@ class itself. ------- str """ + ... + + +class Extension(ExtensionLike): + """ + Abstract base class defining an extension to ASDF. + + Implementing classes must provide the `extension_uri`. + Other properties are optional. + """ + + @classmethod + def __subclasshook__(cls, class_): + if cls is Extension: + return hasattr(class_, "extension_uri") + return NotImplemented # pragma: no cover @property - def legacy_class_names(self): + def legacy_class_names(self) -> Iterable[str]: """ Get the set of fully-qualified class names used by older versions of this extension. This allows a new-style @@ -54,7 +68,7 @@ def legacy_class_names(self): return set() @property - def asdf_standard_requirement(self): + def asdf_standard_requirement(self) -> str | None: """ Get the ASDF Standard version requirement for this extension. @@ -64,10 +78,10 @@ def asdf_standard_requirement(self): If str, PEP 440 version specifier. If None, support all versions. """ - return + return None @property - def converters(self): + def converters(self) -> Iterable[Converter]: """ Get the `asdf.extension.Converter` instances for tags and Python types supported by this extension. @@ -79,7 +93,7 @@ def converters(self): return [] @property - def tags(self): + def tags(self) -> Iterable[str | TagDefinition]: """ Get the YAML tags supported by this extension. @@ -90,7 +104,7 @@ def tags(self): return [] @property - def compressors(self): + def compressors(self) -> Iterable[Compressor]: """ Get the `asdf.extension.Compressor` instances for compression schemes supported by this extension. @@ -102,7 +116,7 @@ def compressors(self): return [] @property - def yaml_tag_handles(self): + def yaml_tag_handles(self) -> dict[str, str]: """ Get a dictionary of custom yaml TAG handles defined by the extension. @@ -120,7 +134,7 @@ def yaml_tag_handles(self): return {} @property - def validators(self): + def validators(self) -> Iterable[Validator]: """ Get the `asdf.extension.Validator` instances for additional schema properties supported by this extension. @@ -132,7 +146,7 @@ def validators(self): return [] -class ExtensionProxy(Extension): +class ExtensionProxy(ExtensionLike): """ Proxy that wraps an extension, provides default implementations of optional methods, and carries additional information on the @@ -140,14 +154,14 @@ class ExtensionProxy(Extension): """ @classmethod - def maybe_wrap(cls, delegate) -> ExtensionProxy: + def maybe_wrap(cls, delegate: ExtensionLike) -> ExtensionProxy: if isinstance(delegate, ExtensionProxy): return delegate return ExtensionProxy(delegate) - def __init__(self, delegate, package_name=None, package_version=None): - if not isinstance(delegate, Extension): + def __init__(self, delegate: ExtensionLike, package_name=None, package_version=None): + if not isinstance(delegate, ExtensionLike): msg = "Extension must implement the Extension interface" raise TypeError(msg) @@ -229,7 +243,7 @@ class itself. return getattr(self._delegate, "extension_uri", None) @property - def legacy_class_names(self): + def legacy_class_names(self) -> set[str]: """ Get the set of fully-qualified class names used by older versions of this extension. This allows a new-style @@ -243,7 +257,7 @@ def legacy_class_names(self): return self._legacy_class_names @property - def asdf_standard_requirement(self): + def asdf_standard_requirement(self) -> SpecifierSet: """ Get the extension's ASDF Standard requirement. @@ -254,7 +268,7 @@ def asdf_standard_requirement(self): return self._asdf_standard_requirement @property - def converters(self): + def converters(self) -> list[ConverterProxy]: """ Get the extension's converters. @@ -265,7 +279,7 @@ def converters(self): return self._converters @property - def compressors(self): + def compressors(self) -> list[Compressor]: """ Get the extension's compressors. @@ -276,7 +290,7 @@ def compressors(self): return self._compressors @property - def tags(self): + def tags(self) -> list[TagDefinition]: """ Get the YAML tags supported by this extension. @@ -287,7 +301,7 @@ def tags(self): return self._tags @property - def types(self): + def types(self) -> list[str | type[Any]]: """ Get the legacy extension's ExtensionType subclasses. @@ -320,7 +334,7 @@ def url_mapping(self): return getattr(self._delegate, "url_mapping", []) @property - def delegate(self): + def delegate(self) -> ExtensionLike: """ Get the wrapped extension instance. @@ -331,7 +345,7 @@ def delegate(self): return self._delegate @property - def package_name(self): + def package_name(self) -> str | None: """ Get the name of the Python package that provided this extension. @@ -343,7 +357,7 @@ def package_name(self): return self._package_name @property - def package_version(self): + def package_version(self) -> str | None: """ Get the version of the Python package that provided the extension @@ -355,7 +369,7 @@ def package_version(self): return self._package_version @property - def class_name(self): + def class_name(self) -> str: """ Get the fully qualified class name of the extension. @@ -366,14 +380,14 @@ def class_name(self): return self._class_name @property - def legacy(self): + def legacy(self) -> bool: """ False """ return self._legacy @property - def yaml_tag_handles(self): + def yaml_tag_handles(self) -> dict[str, str]: """ Get a dictionary of custom yaml TAG handles defined by the extension. @@ -391,7 +405,7 @@ def yaml_tag_handles(self): return self._yaml_tag_handles @property - def validators(self): + def validators(self) -> list[Validator]: """ Get the `asdf.extension.Validator` instances for additional schema properties supported by this extension. diff --git a/asdf/extension/_manager.py b/asdf/extension/_manager.py index 74e8ec2ed..9c0297c92 100644 --- a/asdf/extension/_manager.py +++ b/asdf/extension/_manager.py @@ -5,6 +5,8 @@ from functools import lru_cache from typing import TYPE_CHECKING +from typing_extensions import TypeVar + from asdf.tagged import Tagged from asdf.util import get_class_name, uri_match @@ -15,9 +17,11 @@ from typing import Any from asdf.exceptions import ValidationError - from asdf.extension import Validator + from asdf.extension import Converter, Extension, TagDefinition, Validator from asdf.typing import TreeKey +_T_contra = TypeVar("_T_contra", contravariant=True) + def _resolve_type(path): """ @@ -67,7 +71,7 @@ class ExtensionManager: in the list take precedence. """ - def __init__(self, extensions): + def __init__(self, extensions: Iterable[Extension | ExtensionProxy]): self._extensions = [ExtensionProxy.maybe_wrap(e) for e in extensions] self._tag_defs_by_tag = {} @@ -138,7 +142,7 @@ def __init__(self, extensions): self._validator_manager = _get_cached_validator_manager(tuple(validators)) @property - def extensions(self): + def extensions(self) -> list[ExtensionProxy]: """ Get the list of extensions. @@ -148,7 +152,7 @@ def extensions(self): """ return self._extensions - def handles_tag(self, tag): + def handles_tag(self, tag: str) -> bool: """ Return `True` if the specified tag is handled by a converter. @@ -164,7 +168,7 @@ def handles_tag(self, tag): """ return tag in self._converters_by_tag - def handles_type(self, typ): + def handles_type(self, typ: type[Any]) -> bool: """ Returns `True` if the specified Python type is handled by a converter. @@ -182,7 +186,7 @@ def handles_type(self, typ): self._index_converters() return typ in self._converters_by_type - def handles_tag_definition(self, tag): + def handles_tag_definition(self, tag: str) -> bool: """ Return `True` if the specified tag has a definition. @@ -197,7 +201,7 @@ def handles_tag_definition(self, tag): """ return tag in self._tag_defs_by_tag - def get_tag_definition(self, tag): + def get_tag_definition(self, tag: str) -> TagDefinition: """ Get the tag definition for the specified tag. @@ -221,7 +225,7 @@ def get_tag_definition(self, tag): msg = f"No support available for YAML tag '{tag}'. You may need to install a missing extension." raise KeyError(msg) from None - def get_converter_for_tag(self, tag): + def get_converter_for_tag(self, tag: str) -> Converter[Any]: """ Get the converter for the specified tag. @@ -245,7 +249,7 @@ def get_converter_for_tag(self, tag): msg = f"No support available for YAML tag '{tag}'. You may need to install a missing extension." raise KeyError(msg) from None - def get_converter_for_type(self, typ): + def get_converter_for_type(self, typ: type[_T_contra]) -> Converter[_T_contra]: """ Get the converter for the specified Python type. @@ -273,7 +277,7 @@ def get_converter_for_type(self, typ): ) raise KeyError(msg) from None - def _index_converters(self): + def _index_converters(self) -> None: """ Search _converters_by_class_path for paths (strings) that refer to classes that are currently imported. For imported @@ -290,11 +294,11 @@ def _index_converters(self): del self._converters_by_class_path[class_path] @property - def validator_manager(self): + def validator_manager(self) -> ValidatorManager: return self._validator_manager -def get_cached_extension_manager(extensions): +def get_cached_extension_manager(extensions: Iterable[Extension | ExtensionProxy]) -> ExtensionManager: """ Get a previously created ExtensionManager for the specified extensions, or create and cache one if necessary. Building @@ -323,7 +327,7 @@ def get_cached_extension_manager(extensions): @lru_cache -def _get_cached_extension_manager(extensions): +def _get_cached_extension_manager(extensions: tuple[Extension | ExtensionProxy, ...]) -> ExtensionManager: return ExtensionManager(extensions) @@ -396,7 +400,7 @@ def __call__( yield from validator.validate(schema_property_value, node, schema) -def _validator_matches(validator, node): +def _validator_matches(validator: Validator, node: Any) -> bool: if any(t == "**" for t in validator.tags): return True @@ -407,5 +411,5 @@ def _validator_matches(validator, node): @lru_cache -def _get_cached_validator_manager(validators): +def _get_cached_validator_manager(validators: Iterable[Validator]) -> ValidatorManager: return ValidatorManager(validators) diff --git a/asdf/extension/_manifest.py b/asdf/extension/_manifest.py index 43b5b3610..1b5748b9f 100644 --- a/asdf/extension/_manifest.py +++ b/asdf/extension/_manifest.py @@ -1,8 +1,18 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + import yaml from ._extension import Extension from ._tag import TagDefinition +if TYPE_CHECKING: + from collections.abc import Mapping + + from asdf.extension import Compressor, Converter, Validator + from asdf.typing import TreeKey + class ManifestExtension(Extension): """ @@ -28,7 +38,7 @@ class ManifestExtension(Extension): """ @classmethod - def from_uri(cls, manifest_uri, **kwargs): + def from_uri(cls, manifest_uri: str, **kwargs) -> ManifestExtension: """ Construct the extension using the manifest with the specified URI. The manifest document must be registered @@ -46,7 +56,15 @@ def from_uri(cls, manifest_uri, **kwargs): manifest = yaml.safe_load(get_config().resource_manager[manifest_uri]) return cls(manifest, **kwargs) - def __init__(self, manifest, *, legacy_class_names=None, converters=None, compressors=None, validators=None): + def __init__( + self, + manifest: Mapping[TreeKey, Any], + *, + legacy_class_names: list[str] | None = None, + converters: list[Converter[Any]] | None = None, + compressors: list[Compressor] | None = None, + validators: list[Validator] | None = None, + ): self._manifest = manifest if legacy_class_names is None: @@ -70,15 +88,15 @@ def __init__(self, manifest, *, legacy_class_names=None, converters=None, compre self._validators = validators @property - def extension_uri(self): + def extension_uri(self) -> str: return self._manifest["extension_uri"] @property - def legacy_class_names(self): + def legacy_class_names(self) -> list[str]: return self._legacy_class_names @property - def asdf_standard_requirement(self): + def asdf_standard_requirement(self) -> str | None: version = self._manifest.get("asdf_standard_requirement", None) if version is None: return None @@ -94,20 +112,20 @@ def asdf_standard_requirement(self): return ",".join(specifiers) @property - def converters(self): + def converters(self) -> list[Converter[Any]]: return self._converters @property - def compressors(self): + def compressors(self) -> list[Compressor]: return self._compressors @property - def validators(self): + def validators(self) -> list[Validator]: return self._validators @property - def tags(self): - result = [] + def tags(self) -> list[str | TagDefinition]: + result: list[str | TagDefinition] = [] for tag in self._manifest.get("tags", []): if isinstance(tag, str): # ExtensionProxy knows how to handle str tags. diff --git a/asdf/extension/_serialization_context.py b/asdf/extension/_serialization_context.py index 2630b0a0d..3dff3c209 100644 --- a/asdf/extension/_serialization_context.py +++ b/asdf/extension/_serialization_context.py @@ -1,9 +1,19 @@ +from __future__ import annotations + import enum +from typing import TYPE_CHECKING, Any from asdf._block.key import Key as BlockKey from asdf._helpers import validate_version from asdf.extension._extension import ExtensionProxy +if TYPE_CHECKING: + from asdf import AsdfFile + from asdf._block.manager import Manager as BlockManager + from asdf.extension import ExtensionLike, ExtensionManager + from asdf.typing import ArrayStorage, BlockDataCallback, Compression, NDArray + from asdf.versioning import AsdfVersion + class SerializationContext: """ @@ -14,7 +24,13 @@ class SerializationContext: classes (like Converters) via method arguments. """ - def __init__(self, version, extension_manager, url, blocks): + def __init__( + self, + version: str | AsdfVersion, + extension_manager: ExtensionManager, + url: str | None, + blocks: BlockManager, + ): self._version = validate_version(version) self._extension_manager = extension_manager self._url = url @@ -24,7 +40,7 @@ def __init__(self, version, extension_manager, url, blocks): self.__extensions_used = set() @property - def url(self): + def url(self) -> str | None: """ The URL (if any) of the file being read or written. @@ -39,7 +55,7 @@ def url(self): return self._url @property - def version(self): + def version(self) -> str: """ Get the ASDF Standard version. @@ -50,7 +66,7 @@ def version(self): return self._version @property - def extension_manager(self): + def extension_manager(self) -> ExtensionManager: """ Get the ExtensionManager for enabled extensions. @@ -60,7 +76,7 @@ def extension_manager(self): """ return self._extension_manager - def _mark_extension_used(self, extension): + def _mark_extension_used(self, extension: ExtensionLike) -> None: """ Note that an extension was used when reading or writing the file. @@ -71,7 +87,7 @@ def _mark_extension_used(self, extension): self.__extensions_used.add(ExtensionProxy.maybe_wrap(extension)) @property - def _extensions_used(self): + def _extensions_used(self) -> set[ExtensionProxy]: """ Get the set of extensions that were used when reading or writing the file. @@ -81,7 +97,7 @@ def _extensions_used(self): """ return self.__extensions_used - def get_block_data_callback(self, index, key=None): + def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> BlockDataCallback: """ Generate a callable that when called will read data from an ASDF block at the provided index. @@ -127,7 +143,7 @@ def find_available_block_index(self, data_callback, key=None): """ raise NotImplementedError("abstract") - def generate_block_key(self): + def generate_block_key(self) -> BlockKey: """ Generate a BlockKey used for Converters that wish to use multiple blocks @@ -141,13 +157,13 @@ def generate_block_key(self): """ raise NotImplementedError("abstract") - def assign_object(self, obj): + def assign_object(self, obj: Any) -> None: self._obj = obj - def assign_blocks(self): + def assign_blocks(self) -> None: pass - def set_array_storage(self, arr, array_storage): + def set_array_storage(self, arr: NDArray, array_storage: ArrayStorage) -> None: """ Set the block type to use for the given array data. @@ -171,7 +187,7 @@ def set_array_storage(self, arr, array_storage): """ self._blocks._set_array_storage(arr, array_storage) - def get_array_storage(self, arr): + def get_array_storage(self, arr: NDArray) -> ArrayStorage: """ Get the block type for the given array data. @@ -181,7 +197,7 @@ def get_array_storage(self, arr): """ return self._blocks._get_array_storage(arr) - def set_array_compression(self, arr, compression, **compression_kwargs): + def set_array_compression(self, arr: NDArray, compression: Compression, **compression_kwargs) -> None: """ Set the compression to use for the given array data. @@ -209,7 +225,7 @@ def set_array_compression(self, arr, compression, **compression_kwargs): """ self._blocks._set_array_compression(arr, compression, **compression_kwargs) - def get_array_compression(self, arr): + def get_array_compression(self, arr: NDArray) -> Compression: """ Get the compression type for the given array data. @@ -223,11 +239,11 @@ def get_array_compression(self, arr): """ return self._blocks._get_array_compression(arr) - def get_array_compression_kwargs(self, arr): + def get_array_compression_kwargs(self, arr: NDArray) -> dict[str, Any]: """ """ return self._blocks._get_array_compression_kwargs(arr) - def set_array_save_base(self, arr, save_base): + def set_array_save_base(self, arr: NDArray, save_base: bool | None) -> None: """ Set the ``save_base`` option for ``arr``. When ``arr`` is written to a file, if ``save_base`` is ``True`` the base array @@ -246,7 +262,7 @@ def set_array_save_base(self, arr, save_base): """ self._blocks._set_array_save_base(arr, save_base) - def get_array_save_base(self, arr): + def get_array_save_base(self, arr: NDArray) -> bool | None: """ Returns the ``save_base`` option for ``arr``. When ``arr`` is written to a file, if ``save_base`` is ``True`` the base array @@ -279,13 +295,13 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.assign_object(None) - def assign_object(self, obj): + def assign_object(self, obj: Any) -> None: super().assign_object(obj) if obj is None: self._cb = None self._keys_to_assign = {} - def assign_blocks(self): + def assign_blocks(self) -> None: super().assign_blocks() if self._cb is not None: self._blocks._data_callbacks.assign_object(self._obj, self._cb) @@ -302,7 +318,7 @@ def assign_blocks(self): # assigned object self.assign_object(None) - def get_block_data_callback(self, index, key=None): + def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> BlockDataCallback: if key is None: if self._cb is not None: # this operation has already accessed a block without using @@ -322,7 +338,7 @@ def get_block_data_callback(self, index, key=None): self._keys_to_assign[key] = cb return cb - def generate_block_key(self): + def generate_block_key(self) -> BlockKey: key = BlockKey() self._keys_to_assign[key] = None return key @@ -339,12 +355,12 @@ class WriteBlocksContext(SerializationContext): being serialized. """ - def find_available_block_index(self, data_callback, key=None): + def find_available_block_index(self, data_callback: BlockDataCallback, key: BlockKey | None = None) -> int | str: if key is None: key = self._obj return self._blocks.make_write_block(data_callback, None, key) - def generate_block_key(self): + def generate_block_key(self) -> BlockKey: return BlockKey(self._obj) @@ -359,7 +375,7 @@ class BlockAccess(enum.Enum): READ = ReadBlocksContext -def create(asdf_file, block_access=BlockAccess.NONE): +def create(asdf_file: AsdfFile, block_access: BlockAccess = BlockAccess.NONE) -> SerializationContext: """ Create a SerializationContext instance (or subclass) using an AsdfFile instance, asdf_file. diff --git a/asdf/extension/_tag.py b/asdf/extension/_tag.py index e96c4e67f..ccc73033d 100644 --- a/asdf/extension/_tag.py +++ b/asdf/extension/_tag.py @@ -15,7 +15,14 @@ class TagDefinition: Long description of the tag. """ - def __init__(self, tag_uri, *, schema_uris=None, title=None, description=None): + def __init__( + self, + tag_uri: str, + *, + schema_uris: str | None = None, + title: str | None = None, + description: str | None = None, + ): if "*" in tag_uri: msg = "URI patterns are not permitted in TagDefinition" raise ValueError(msg) @@ -35,7 +42,7 @@ def __init__(self, tag_uri, *, schema_uris=None, title=None, description=None): self._description = description @property - def tag_uri(self): + def tag_uri(self) -> str: """ Get the tag URI. @@ -46,7 +53,7 @@ def tag_uri(self): return self._tag_uri @property - def schema_uris(self): + def schema_uris(self) -> list[str]: """ Get the URIs of the schemas that should be used to validate objects with this tag. @@ -58,7 +65,7 @@ def schema_uris(self): return self._schema_uris @property - def title(self): + def title(self) -> str | None: """ Get the short description of the tag. @@ -69,7 +76,7 @@ def title(self): return self._title @property - def description(self): + def description(self) -> str | None: """ Get the long description of the tag. @@ -79,5 +86,5 @@ def description(self): """ return self._description - def __repr__(self): + def __repr__(self) -> str: return f"" diff --git a/asdf/extension/_validator.py b/asdf/extension/_validator.py index 155e2d3b3..36a437260 100644 --- a/asdf/extension/_validator.py +++ b/asdf/extension/_validator.py @@ -17,14 +17,16 @@ class Validator(abc.ABC): in ASDF schemas. """ - @abc.abstractproperty + @property + @abc.abstractmethod def schema_property(self) -> str: """ Name of the schema property used to invoke this validator. """ ... - @abc.abstractproperty + @property + @abc.abstractmethod def tags(self) -> Iterable[str]: """ Get the YAML tags that are appropriate to this validator. diff --git a/asdf/tags/core/integer.py b/asdf/tags/core/integer.py index 6f5968fce..6da744c89 100644 --- a/asdf/tags/core/integer.py +++ b/asdf/tags/core/integer.py @@ -1,4 +1,16 @@ +from __future__ import annotations + from numbers import Integral +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import numpy as np + + from asdf.typing import ArrayStorage + + # Using a custom definition because Integral doesn't work with type-checking + # and numbers._IntegralLike doesn't support comparison operations + _IntegerLike = int | np.integer[Any] class IntegerType: @@ -38,13 +50,13 @@ class IntegerType: ... assert aa["largeval"] == largeval """ - def __init__(self, value, storage_type="internal"): + def __init__(self, value: _IntegerLike, storage_type: ArrayStorage = "internal"): if storage_type not in ["internal", "inline"]: msg = f"storage_type '{storage_type}' is not a recognized storage type" raise ValueError(msg) - self._value = value - self._sign = "-" if value < 0 else "+" - self._storage = storage_type + self._value: _IntegerLike = value + self._sign: str = "-" if value < 0 else "+" + self._storage: ArrayStorage = storage_type def __int__(self): return int(self._value) diff --git a/asdf/tags/core/ndarray.py b/asdf/tags/core/ndarray.py index 03577743f..432dab671 100644 --- a/asdf/tags/core/ndarray.py +++ b/asdf/tags/core/ndarray.py @@ -3,6 +3,7 @@ import mmap import sys import typing +from typing import Any import numpy as np from numpy import ma @@ -11,6 +12,8 @@ from asdf._jsonschema import ValidationError if typing.TYPE_CHECKING: + import numpy.typing as npt + from asdf.typing import NDArray _STRUCTURED_DATATYPE_KEYS = {"name", "datatype", "byteorder", "shape"} @@ -283,7 +286,7 @@ def __init__(self, source, shape, dtype, offset, strides, order, mask, data_call self._strides = strides self._order = order - def _make_array(self): + def _make_array(self) -> npt.NDArray[Any]: # If the ASDF file has been updated in-place, then there's # a chance that the block's original data object has been # closed and replaced. We need to check here and re-generate @@ -303,29 +306,32 @@ def _make_array(self): self._array = None del fd - if self._array is None: - if isinstance(self._source, str): - # we need to keep _source as a str to allow stdatamodels to - # support AsdfInFits - data = self._data_callback() - else: - # cached data is used here so that multiple NDArrayTypes will all use - # the same base array - data = self._data_callback(_attr="cached_data") - - if hasattr(data, "base") and isinstance(data.base, mmap.mmap) and data.base.closed: - msg = "ASDF file has already been closed. Can not get the data." - raise OSError(msg) - - # compute shape (streaming blocks have '0' data size in the block header) - shape = self.get_actual_shape( - self._shape, - self._strides, - self._dtype, - data.size, - ) - self._array = np.ndarray(shape, self._dtype, data, self._offset, self._strides, self._order) - self._array = self._apply_mask(self._array, self._mask) + if self._array is not None: + return self._array + + if isinstance(self._source, str): + # we need to keep _source as a str to allow stdatamodels to + # support AsdfInFits + data = self._data_callback() + else: + # cached data is used here so that multiple NDArrayTypes will all use + # the same base array + data = self._data_callback(_attr="cached_data") + + if hasattr(data, "base") and isinstance(data.base, mmap.mmap) and data.base.closed: + msg = "ASDF file has already been closed. Can not get the data." + raise OSError(msg) + + # compute shape (streaming blocks have '0' data size in the block header) + shape = self.get_actual_shape( + self._shape, + self._strides, + self._dtype, + data.size, + ) + self._array = np.ndarray(shape, self._dtype, data, self._offset, self._strides, self._order) + self._array = self._apply_mask(self._array, self._mask) + return self._array def _apply_mask(self, array, mask): diff --git a/asdf/tags/core/stream.py b/asdf/tags/core/stream.py index fed1c84c8..1b1e161d5 100644 --- a/asdf/tags/core/stream.py +++ b/asdf/tags/core/stream.py @@ -34,4 +34,4 @@ def __repr__(self): return f"Stream({self._shape}, {self._datatype}, strides={self._strides})" def __str__(self): - return str(self.__repr__()) + return self.__repr__() diff --git a/asdf/typing.py b/asdf/typing.py index 2a9772caa..b648113ba 100644 --- a/asdf/typing.py +++ b/asdf/typing.py @@ -2,15 +2,16 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from pathlib import Path -from typing import Any, Literal, Protocol, TypeAlias +from typing import Any, Literal, TypeAlias import numpy as np import numpy.typing as npt -from typing_extensions import Reader, Writer +from typing_extensions import Reader, TypeVar, Writer from asdf.generic_io import GenericFile +from asdf.tags.core import NDArrayType from asdf.util import _NOT_SET_TYPE from asdf.versioning import AsdfVersion @@ -20,7 +21,6 @@ "BlockDataCallback", "ByteArray1D", "Compression", - "ExtensionLike", "FileLike", "FileMode", "FilterFn", @@ -33,15 +33,6 @@ ] -# Alternate version of `Extension` for use in type-hints -# The way `Extension` works is weird enough that it can't be replaced in actual code without a lot of changes -class ExtensionLike(Protocol): - """Object that contains an extension URI and can be wrapped by ``ExtensionProxy``.""" - - @property - def extension_uri(self) -> str | None: ... - - # Ideally this would be `str | int | bool` # Unfortunately this becomes a headache since mapping keys aren't covariant # See: https://github.com/python/typing/pull/273 @@ -51,13 +42,15 @@ def extension_uri(self) -> str | None: ... #: Valid ASDF tree keys TreeKey: TypeAlias = Any +#: A YAML node that can be passed to or returned from an ASDF converter +YamlNode: TypeAlias = Mapping[TreeKey, Any] | Sequence[Any] | str #: Local file path or remote file URI PathLike: TypeAlias = str | Path #: Readable/writable file object or the path or URI of an openable file FileLike: TypeAlias = PathLike | Reader | Writer | GenericFile #: A type interpretable as a version number -AsdfVersionLike: TypeAlias = AsdfVersion | str | list[int] | tuple[int, ...] +AsdfVersionLike: TypeAlias = AsdfVersion | str #: Supported modes for opening a file FileMode: TypeAlias = Literal["r", "w", "rw"] @@ -72,12 +65,16 @@ def extension_uri(self) -> str | None: ... FilterFn: TypeAlias = Callable[[Any], bool] | Callable[[Any, Any], bool] #: ASDF-compatible multi-dimensional array -NDArray: TypeAlias = npt.NDArray[Any] +NDArray: TypeAlias = npt.NDArray[Any] | NDArrayType #: A 1-D byte numpy array used to read and write block data ByteArray1D: TypeAlias = np.ndarray[tuple[int], np.dtype[np.uint8]] +_Array = TypeVar("_Array", default=NDArray) + +#: A callback that returns a numpy array +ArrayCallback = Callable[[], _Array] #: A callback that returns a `ByteArray1D` -BlockDataCallback = Callable[[], ByteArray1D] +BlockDataCallback = ArrayCallback[ByteArray1D] NotSetType = Literal[_NOT_SET_TYPE.NOT_SET] diff --git a/asdf/versioning.py b/asdf/versioning.py index 86cb0b053..8932dc319 100644 --- a/asdf/versioning.py +++ b/asdf/versioning.py @@ -3,6 +3,8 @@ of the ASDF spec. """ +from __future__ import annotations + from functools import total_ordering import yaml @@ -22,7 +24,7 @@ def get_supported_core_schema_versions(): __all__ = ["AsdfVersion", "AsdfVersionMixin", "join_tag_version", "split_tag_version"] -def split_tag_version(tag): +def split_tag_version(tag: str) -> tuple[str, AsdfVersion]: """ Split a tag into its base and version. """ @@ -31,7 +33,7 @@ def split_tag_version(tag): return name, version -def join_tag_version(name, version): +def join_tag_version(name: str, version: str) -> str: """ Join the root and version of a tag back together. """ @@ -91,9 +93,11 @@ def __init__(self, version): super().__init__(version) -supported_versions = tuple(AsdfVersion(version) for version in get_supported_core_schema_versions()) +supported_versions: tuple[AsdfVersion, ...] = tuple( + AsdfVersion(version) for version in get_supported_core_schema_versions() +) -default_version = supported_versions[-1] +default_version: AsdfVersion = supported_versions[-1] # This is the ASDF core schemas version at which the format of the history # field changed to include extension metadata. From 7144b307fa8cef484c5dc9997ca60567bb9db92e Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Mon, 10 Aug 2026 10:31:25 -0400 Subject: [PATCH 3/8] Added type hints to `_node_info.py` --- asdf/_node_info.py | 87 ++++++++++++++++++++++++++++------------------ 1 file changed, 53 insertions(+), 34 deletions(-) diff --git a/asdf/_node_info.py b/asdf/_node_info.py index 87634f131..a0e26c30e 100644 --- a/asdf/_node_info.py +++ b/asdf/_node_info.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from typing_extensions import NamedTuple @@ -9,7 +10,10 @@ from .treeutil import get_children, is_container if TYPE_CHECKING: - from asdf.typing import TreeKey + from collections.abc import Collection, Mapping + + from asdf.extension import ExtensionManager + from asdf.typing import FilterFn, TreeKey def _filter_tree(info, filters): @@ -145,7 +149,13 @@ def _get_schema_key(schema, key): return None -def create_tree(key, node, identifier="root", filters=None, extension_manager=None): +def create_tree( + key: str, + node: Mapping[TreeKey, Any], + identifier: str = "root", + filters: Collection[FilterFn] | None = None, + extension_manager: ExtensionManager | None = None, +) -> NodeSchemaInfo | None: """ Create a `NodeSchemaInfo` tree which can be filtered from a base node. @@ -176,13 +186,13 @@ def create_tree(key, node, identifier="root", filters=None, extension_manager=No def collect_schema_info( - key, - path, - node, - identifier="root", - filters=None, - preserve_list=True, - extension_manager=None, + key: str, + path: str | None, + node: Any, + identifier: str = "root", + filters: Collection[FilterFn] | None = None, + preserve_list: bool = True, + extension_manager: ExtensionManager | None = None, ): """ Collect from the underlying schemas any of the info stored under key, relative to the path @@ -209,6 +219,8 @@ def collect_schema_info( filters=[] if filters is None else filters, extension_manager=extension_manager, ) + if schema_info is None: + return None info = schema_info.collect_info(preserve_list=preserve_list) @@ -257,6 +269,7 @@ def __repr__(self): return f"{self.info}" +@dataclass class NodeSchemaInfo: """ Container for keyed information collected from a schema about a node of an ASDF file tree. @@ -302,24 +315,23 @@ class NodeSchemaInfo: The portion of the underlying schema corresponding to the node. """ - def __init__(self, key, parent, identifier, node, depth, recursive=False, visible=True, extension_manager=None): - self.key = key - self.parent = parent - self.identifier = identifier - self.node = node - self.depth = depth - self.recursive = recursive - self.visible = visible - self.children = [] - self.schema = None - self.extension_manager = extension_manager or _get_extension_manager() + key: str + parent: NodeSchemaInfo | None + identifier: str + node: Any + depth: int + recursive: bool = False + visible: bool = True + extension_manager: ExtensionManager = field(default_factory=_get_extension_manager) + children: list[NodeSchemaInfo] = field(default_factory=list, init=False) + schema: Mapping[TreeKey, Any] | None = field(default=None, init=False) @property - def visible_children(self): + def visible_children(self) -> list[NodeSchemaInfo]: return [c for c in self.children if c.visible] @property - def parent_node(self): + def parent_node(self) -> Any | None: if self.parent is not None: return self.parent.node @@ -349,7 +361,14 @@ def set_schema_from_node(self, node, extension_manager): self.schema = schema @classmethod - def from_root_node(cls, key, root_identifier, root_node, schema=None, extension_manager=None): + def from_root_node( + cls, + key: str, + root_identifier: str, + root_node: Any, + schema: Mapping[TreeKey, Any] | None = None, + extension_manager: ExtensionManager | None = None, + ): """ Build a NodeSchemaInfo tree from the given ASDF root node. Intentionally processes the tree in breadth-first order so that recursively @@ -431,7 +450,7 @@ def from_root_node(cls, key, root_identifier, root_node, schema=None, extension_ return root_info - def collect_info(self, preserve_list=True): + def collect_info(self, preserve_list: bool = True) -> dict[str, Any] | list[Any]: """ Collect the information from the NodeSchemaInfo tree, and return it as nested dict. @@ -442,15 +461,15 @@ def collect_info(self, preserve_list=True): If True, then lists are preserved. Otherwise, they are turned into dicts. """ if preserve_list and isinstance(self.node, (list, tuple)) and self.info is None: - info = [c_info for child in self.visible_children if len(c_info := child.collect_info(preserve_list)) > 0] - else: - info = { - child.identifier: c_info - for child in self.visible_children - if len(c_info := child.collect_info(preserve_list)) > 0 - } - - if self.info is not None: - info[self.key] = SchemaInfo(self.info, self.node) + return [c_info for child in self.visible_children if len(c_info := child.collect_info(preserve_list)) > 0] + + info: dict[str, Any] = { + child.identifier: c_info + for child in self.visible_children + if len(c_info := child.collect_info(preserve_list)) > 0 + } + + if self.info is not None: + info[self.key] = SchemaInfo(self.info, self.node) return info From c70100e7bb1e99e69a3c11953dcb655293b85241 Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Mon, 10 Aug 2026 10:32:19 -0400 Subject: [PATCH 4/8] Fixed test suite typing errors --- asdf/_asdf.py | 2 +- asdf/_tests/tags/core/tests/test_integer.py | 10 +++- asdf/_tests/test_config.py | 2 +- asdf/_tests/test_extension.py | 55 ++++++++++++++------- asdf/_tests/test_info.py | 8 +-- asdf/_tests/test_resource.py | 2 +- asdf/_tests/test_serialization_context.py | 9 ++-- asdf/extension/_manager.py | 12 ++--- asdf/extension/_tag.py | 2 +- asdf/resource.py | 24 ++++----- 10 files changed, 77 insertions(+), 49 deletions(-) diff --git a/asdf/_asdf.py b/asdf/_asdf.py index 5f1bfbba5..5b8151d71 100644 --- a/asdf/_asdf.py +++ b/asdf/_asdf.py @@ -328,7 +328,7 @@ def _check_extensions(self, tree: Mapping[TreeKey, Any], strict: bool = False) - package_description = f"{package_name}=={package_version}" installed_version = None for mapping in get_config().resource_manager._resource_mappings: - if mapping.package_name == package_name: + if mapping.package_name == package_name and mapping.package_version is not None: installed_version = Version(mapping.package_version) break msg = None diff --git a/asdf/_tests/tags/core/tests/test_integer.py b/asdf/_tests/tags/core/tests/test_integer.py index ee4e35129..9247c6433 100644 --- a/asdf/_tests/tags/core/tests/test_integer.py +++ b/asdf/_tests/tags/core/tests/test_integer.py @@ -1,11 +1,16 @@ import random +from typing import TYPE_CHECKING import pytest import asdf +import asdf.util from asdf import IntegerType from asdf.testing.helpers import roundtrip_object +if TYPE_CHECKING: + from asdf.typing import ArrayStorage + # Make sure tests are deterministic random.seed(0) @@ -34,7 +39,8 @@ def test_integer_value(value, sign): def test_integer_storage(tmp_path, inline): tmpfile = str(tmp_path / "integer.asdf") - kwargs = {} + # This typing is a little hacky but whatever + kwargs: dict[str, ArrayStorage] = {} if inline: kwargs["storage_type"] = "inline" @@ -63,5 +69,5 @@ def test_integer_conversion(): integer = asdf.IntegerType(value) assert integer == value - assert int(integer) == int(value) + assert int(integer) == value assert float(integer) == float(value) diff --git a/asdf/_tests/test_config.py b/asdf/_tests/test_config.py index 08bf088c1..ea4f22832 100644 --- a/asdf/_tests/test_config.py +++ b/asdf/_tests/test_config.py @@ -37,7 +37,7 @@ def test_config_context_nested(): def test_config_context_threaded(): assert get_config().validate_on_read is True - thread_value = None + thread_value: bool | None = None def worker(): nonlocal thread_value diff --git a/asdf/_tests/test_extension.py b/asdf/_tests/test_extension.py index 28f835aa8..3d179d5b9 100644 --- a/asdf/_tests/test_extension.py +++ b/asdf/_tests/test_extension.py @@ -1,6 +1,8 @@ import collections import fractions import sys +import typing +from typing import Any import pytest from packaging.specifiers import SpecifierSet @@ -16,6 +18,7 @@ ExtensionManager, ExtensionProxy, ManifestExtension, + SerializationContext, TagDefinition, Validator, get_cached_extension_manager, @@ -90,17 +93,17 @@ def __init__(self, tags=None, types=None): self._types = types @property - def tags(self): + def tags(self) -> list[str]: return self._tags @property - def types(self): + def types(self) -> list[str | type]: return self._types - def to_yaml_tree(self, obj, tag, ctx): + def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> str: return "to_yaml_tree result" - def from_yaml_tree(self, obj, tag, ctx): + def from_yaml_tree(self, node: str, tag: str, ctx: SerializationContext) -> Any: return "from_yaml_tree result" @@ -152,6 +155,7 @@ def test_extension_proxy_maybe_wrap(): assert ExtensionProxy.maybe_wrap(proxy) is proxy with pytest.raises(TypeError, match=r"Extension must implement the Extension interface"): + # pyrefly: ignore [bad-argument-type] ExtensionProxy.maybe_wrap(object()) @@ -230,6 +234,7 @@ def test_extension_proxy(): # Should fail when the input is not one of the two extension interfaces: with pytest.raises(TypeError, match=r"Extension must implement the Extension interface"): + # pyrefly: ignore [bad-argument-type] ExtensionProxy(object) # Should fail with a bad converter: @@ -467,14 +472,15 @@ def test_converter_proxy(): # Test the minimum set of converter methods: extension = ExtensionProxy(MinimumExtension()) converter = MinimumConverter() + ctx = typing.cast("SerializationContext", None) proxy = ConverterProxy(converter, extension) assert isinstance(proxy, Converter) assert proxy.tags == [] assert proxy.types == [] - assert proxy.to_yaml_tree(None, None, None) == "to_yaml_tree result" - assert proxy.from_yaml_tree(None, None, None) == "from_yaml_tree result" + assert proxy.to_yaml_tree(None, "", ctx) == "to_yaml_tree result" + assert proxy.from_yaml_tree("", "", ctx) == "from_yaml_tree result" assert proxy.tags == [] assert proxy.delegate is converter assert proxy.extension == extension @@ -485,9 +491,12 @@ def test_converter_proxy(): # Check the __eq__ and __hash__ behavior: assert proxy == ConverterProxy(converter, extension) assert proxy != ConverterProxy(MinimumConverter(), extension) - assert proxy != ConverterProxy(converter, MinimumExtension()) + assert proxy != ConverterProxy(converter, ExtensionProxy.maybe_wrap(MinimumExtension())) assert proxy in {ConverterProxy(converter, extension)} - assert proxy not in {ConverterProxy(MinimumConverter(), extension), ConverterProxy(converter, MinimumExtension())} + assert proxy not in { + ConverterProxy(MinimumConverter(), extension), + ConverterProxy(converter, ExtensionProxy.maybe_wrap(MinimumExtension())), + } # Check the __repr__: assert "class: asdf._tests.test_extension.MinimumConverter" in repr(proxy) @@ -525,9 +534,9 @@ def test_converter_proxy(): assert "asdf://somewhere.org/extensions/test/tags/foo-1.0" in proxy.tags assert "asdf://somewhere.org/extensions/test/tags/bar-1.0" in proxy.tags assert proxy.types == [FooType, BarType] - assert proxy.to_yaml_tree(None, None, None) == "to_yaml_tree result" - assert proxy.from_yaml_tree(None, None, None) == "from_yaml_tree result" - assert proxy.select_tag(None, None) == "select_tag result" + assert proxy.to_yaml_tree(None, "", ctx) == "to_yaml_tree result" + assert proxy.from_yaml_tree("", "", ctx) == "from_yaml_tree result" + assert proxy.select_tag(None, ctx) == "select_tag result" assert proxy.delegate is converter assert proxy.extension == extension_proxy assert proxy.package_name == "foo" @@ -540,18 +549,26 @@ def test_converter_proxy(): # Should error because object() does fulfill the Converter interface: with pytest.raises(TypeError, match=r"Converter must implement the .*"): - ConverterProxy(object(), extension) + ConverterProxy( + object(), # pyrefly: ignore [bad-argument-type] + ExtensionProxy.maybe_wrap(extension), + ) # Should fail because tags must be str: with pytest.raises(TypeError, match=r"Converter property .* must contain str values"): - ConverterProxy(MinimumConverter(tags=[object()]), extension) + ConverterProxy( + MinimumConverter(tags=[object()]), + ExtensionProxy.maybe_wrap(extension), + ) # Should fail because types must instances of type: with pytest.raises(TypeError, match=r"Converter property .* must contain str or type values"): # as the code will ignore types if no relevant tags are found # include a tag from this extension to make sure the proxy considers # the types - ConverterProxy(MinimumConverter(tags=[extension.tags[0].tag_uri], types=[object()]), extension) + ConverterProxy( + MinimumConverter(tags=[extension.tags[0].tag_uri], types=[object()]), ExtensionProxy.maybe_wrap(extension) + ) def test_converter_subclass_with_no_supported_tags(): @@ -635,13 +652,13 @@ class FooConverter: tags = ["asdf://somewhere.org/tags/bar", "asdf://somewhere.org/tags/baz"] types = [] - def select_tag(self, *args): + def select_tag(self, obj: Any, ctx: SerializationContext) -> str | None: pass - def to_yaml_tree(self, *args): + def to_yaml_tree(self, obj: Any, tag: str, ctx: SerializationContext) -> Any: pass - def from_yaml_tree(self, *args): + def from_yaml_tree(self, node: Any, tag: str, ctx: SerializationContext) -> Any: pass converter = FooConverter() @@ -897,14 +914,14 @@ def test_warning_or_error_for_default_select_tag(is_subclass, indirect): class Foo: pass - ParentClass = Converter if is_subclass else object + ParentClass: type[Any] = Converter if is_subclass else object if indirect: class IntermediateClass(ParentClass): pass - ParentClass = IntermediateClass + ParentClass: type[Any] = IntermediateClass class FooConverter(ParentClass): tags = ["asdf://somewhere.org/tags/foo-*"] diff --git a/asdf/_tests/test_info.py b/asdf/_tests/test_info.py index 57baa7cfb..affc07f28 100644 --- a/asdf/_tests/test_info.py +++ b/asdf/_tests/test_info.py @@ -8,6 +8,8 @@ import pytest import asdf +import asdf.extension +from asdf._node_info import NodeSchemaInfo from asdf.extension import ExtensionProxy, ManifestExtension from asdf.resource import DirectoryResourceMapping @@ -721,7 +723,7 @@ def __str__(self): ], ) def test_node_property(schema, expected): - ni = asdf._node_info.NodeSchemaInfo.from_root_node("title", "root", {}, schema) + ni = NodeSchemaInfo.from_root_node("title", "root", {}, schema) assert ni.get_schema_for_property("foo") == expected @@ -747,7 +749,7 @@ def test_node_property(schema, expected): ], ) def test_node_property_error(schema): - ni = asdf._node_info.NodeSchemaInfo.from_root_node("title", "root", {}, schema) + ni = NodeSchemaInfo.from_root_node("title", "root", {}, schema) assert ni.get_schema_for_property("foo") == {} @@ -766,7 +768,7 @@ def test_node_property_error(schema): ], ) def test_node_info(schema, expected): - ni = asdf._node_info.NodeSchemaInfo.from_root_node("title", "root", {}, schema) + ni = NodeSchemaInfo.from_root_node("title", "root", {}, schema) assert ni.info == expected diff --git a/asdf/_tests/test_resource.py b/asdf/_tests/test_resource.py index 4d4c325bf..c8fdb7ef6 100644 --- a/asdf/_tests/test_resource.py +++ b/asdf/_tests/test_resource.py @@ -15,7 +15,7 @@ def test_resource_manager(): "http://somewhere.org/schemas/baz-1.0.0": b"baz", "http://somewhere.org/schemas/foz-1.0.0": "foz", } - manager = ResourceManager([mapping1, mapping2]) + manager = ResourceManager([ResourceMappingProxy.maybe_wrap(m) for m in [mapping1, mapping2]]) assert isinstance(manager, Mapping) diff --git a/asdf/_tests/test_serialization_context.py b/asdf/_tests/test_serialization_context.py index 1551e166f..c03760e89 100644 --- a/asdf/_tests/test_serialization_context.py +++ b/asdf/_tests/test_serialization_context.py @@ -3,13 +3,15 @@ import asdf from asdf import get_config +from asdf._block.manager import Manager as BlockManager from asdf.extension import ExtensionManager from asdf.extension._serialization_context import BlockAccess, SerializationContext -def test_serialization_context(): +def test_serialization_context() -> None: extension_manager = ExtensionManager([]) - context = SerializationContext("1.4.0", extension_manager, "file://test.asdf", None) + blocks = BlockManager() + context = SerializationContext("1.4.0", extension_manager, "file://test.asdf", blocks) assert context.version == "1.4.0" assert context.extension_manager is extension_manager assert context._extensions_used == set() @@ -25,10 +27,11 @@ def test_serialization_context(): assert context.url == context._url == "file://test.asdf" with pytest.raises(TypeError, match=r"Extension must implement the Extension interface"): + # pyrefly: ignore [bad-argument-type] context._mark_extension_used(object()) with pytest.raises(ValueError, match=r"ASDF Standard version .* is not supported by asdf==.*"): - SerializationContext("0.5.4", extension_manager, None, None) + SerializationContext("0.5.4", extension_manager, None, blocks) def test_get_block_data_callback(tmp_path): diff --git a/asdf/extension/_manager.py b/asdf/extension/_manager.py index 9c0297c92..fc5a2bee2 100644 --- a/asdf/extension/_manager.py +++ b/asdf/extension/_manager.py @@ -17,7 +17,7 @@ from typing import Any from asdf.exceptions import ValidationError - from asdf.extension import Converter, Extension, TagDefinition, Validator + from asdf.extension import ConverterProxy, Extension, ExtensionLike, TagDefinition, Validator from asdf.typing import TreeKey _T_contra = TypeVar("_T_contra", contravariant=True) @@ -71,7 +71,7 @@ class ExtensionManager: in the list take precedence. """ - def __init__(self, extensions: Iterable[Extension | ExtensionProxy]): + def __init__(self, extensions: Iterable[ExtensionLike]): self._extensions = [ExtensionProxy.maybe_wrap(e) for e in extensions] self._tag_defs_by_tag = {} @@ -225,7 +225,7 @@ def get_tag_definition(self, tag: str) -> TagDefinition: msg = f"No support available for YAML tag '{tag}'. You may need to install a missing extension." raise KeyError(msg) from None - def get_converter_for_tag(self, tag: str) -> Converter[Any]: + def get_converter_for_tag(self, tag: str) -> ConverterProxy[Any]: """ Get the converter for the specified tag. @@ -249,7 +249,7 @@ def get_converter_for_tag(self, tag: str) -> Converter[Any]: msg = f"No support available for YAML tag '{tag}'. You may need to install a missing extension." raise KeyError(msg) from None - def get_converter_for_type(self, typ: type[_T_contra]) -> Converter[_T_contra]: + def get_converter_for_type(self, typ: type[_T_contra]) -> ConverterProxy[_T_contra]: """ Get the converter for the specified Python type. @@ -298,7 +298,7 @@ def validator_manager(self) -> ValidatorManager: return self._validator_manager -def get_cached_extension_manager(extensions: Iterable[Extension | ExtensionProxy]) -> ExtensionManager: +def get_cached_extension_manager(extensions: Iterable[ExtensionLike]) -> ExtensionManager: """ Get a previously created ExtensionManager for the specified extensions, or create and cache one if necessary. Building @@ -313,8 +313,6 @@ def get_cached_extension_manager(extensions: Iterable[Extension | ExtensionProxy ------- asdf.extension.ExtensionManager """ - from ._extension import ExtensionProxy - # The tuple makes the extensions hashable so that we # can pass them to the lru_cache method. The ExtensionProxy # overrides __hash__ to return the hashed object id of the wrapped diff --git a/asdf/extension/_tag.py b/asdf/extension/_tag.py index ccc73033d..e879335dd 100644 --- a/asdf/extension/_tag.py +++ b/asdf/extension/_tag.py @@ -19,7 +19,7 @@ def __init__( self, tag_uri: str, *, - schema_uris: str | None = None, + schema_uris: list[str] | str | None = None, title: str | None = None, description: str | None = None, ): diff --git a/asdf/resource.py b/asdf/resource.py index 8eb1ef919..8dd0672fd 100644 --- a/asdf/resource.py +++ b/asdf/resource.py @@ -4,7 +4,7 @@ """ import pkgutil -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from asdf_standard import DirectoryResourceMapping as _DirectoryResourceMapping @@ -17,6 +17,8 @@ "ResourceMappingProxy", ] +_ResourceMapping = Mapping[str, str | bytes] + class DirectoryResourceMapping(_DirectoryResourceMapping): """ @@ -26,7 +28,7 @@ class DirectoryResourceMapping(_DirectoryResourceMapping): """ -class ResourceMappingProxy(Mapping): +class ResourceMappingProxy(_ResourceMapping): """ Wrapper around a resource mapping that carries additional information on the package that provided @@ -71,7 +73,7 @@ def delegate(self): return self._delegate @property - def package_name(self): + def package_name(self) -> str | None: """ Get the name of the Python package that provided this mapping. @@ -83,7 +85,7 @@ def package_name(self): return self._package_name @property - def package_version(self): + def package_version(self) -> str | None: """ Get the version of the Python package that provided the mapping. @@ -95,7 +97,7 @@ def package_version(self): return self._package_version @property - def class_name(self): + def class_name(self) -> str: """ " Get the fully qualified class name of the mapping. @@ -123,7 +125,7 @@ def __repr__(self): return f"" -class ResourceManager(Mapping): +class ResourceManager(_ResourceMapping): """ Wraps multiple resource mappings into a single interface with some friendlier error handling. @@ -135,16 +137,16 @@ class ResourceManager(Mapping): the first mapping takes precedence. """ - def __init__(self, resource_mappings): + def __init__(self, resource_mappings: Iterable[ResourceMappingProxy]): self._resource_mappings = resource_mappings - self._mappings_by_uri = {} + self._mappings_by_uri: dict[str, ResourceMappingProxy] = {} for mapping in resource_mappings: for uri in mapping: if uri not in self._mappings_by_uri: self._mappings_by_uri[uri] = mapping - def __getitem__(self, uri): + def __getitem__(self, uri: str) -> str | bytes: if uri not in self._mappings_by_uri: msg = f"Resource unavailable for URI: {uri}" raise KeyError(msg) @@ -155,13 +157,13 @@ def __getitem__(self, uri): return content - def __len__(self): + def __len__(self) -> int: return len(self._mappings_by_uri) def __iter__(self): yield from self._mappings_by_uri - def __contains__(self, uri): + def __contains__(self, uri: object) -> bool: # Implement __contains__ only for efficiency. return uri in self._mappings_by_uri From efa6c93fbe16d8b9104226e1e687ea92966f745f Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Tue, 11 Aug 2026 12:01:08 -0400 Subject: [PATCH 5/8] Fixed more typing errors --- asdf/_node_info.py | 11 +++++-- asdf/_tests/tags/core/tests/test_ndarray.py | 4 ++- asdf/_tests/test_config.py | 29 ++++++++++------- asdf/extension/_serialization_context.py | 7 +++-- asdf/tags/core/ndarray.py | 35 +++++++++++++++------ asdf/typing.py | 26 ++++++++++++--- 6 files changed, 79 insertions(+), 33 deletions(-) diff --git a/asdf/_node_info.py b/asdf/_node_info.py index a0e26c30e..ef9a639e5 100644 --- a/asdf/_node_info.py +++ b/asdf/_node_info.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +import typing from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -317,7 +318,7 @@ class NodeSchemaInfo: key: str parent: NodeSchemaInfo | None - identifier: str + identifier: str | int node: Any depth: int recursive: bool = False @@ -450,7 +451,7 @@ def from_root_node( return root_info - def collect_info(self, preserve_list: bool = True) -> dict[str, Any] | list[Any]: + def collect_info(self, preserve_list: bool = True) -> dict[str, Any]: """ Collect the information from the NodeSchemaInfo tree, and return it as nested dict. @@ -460,10 +461,14 @@ def collect_info(self, preserve_list: bool = True) -> dict[str, Any] | list[Any] preserve_list : bool If True, then lists are preserved. Otherwise, they are turned into dicts. """ + # The root node can't be a list but this is hard to encode in the type system + return typing.cast("dict[str, Any]", self._collect_info(preserve_list)) + + def _collect_info(self, preserve_list: bool = True) -> dict[str | int, Any] | list[Any]: if preserve_list and isinstance(self.node, (list, tuple)) and self.info is None: return [c_info for child in self.visible_children if len(c_info := child.collect_info(preserve_list)) > 0] - info: dict[str, Any] = { + info: dict[str | int, Any] = { child.identifier: c_info for child in self.visible_children if len(c_info := child.collect_info(preserve_list)) > 0 diff --git a/asdf/_tests/tags/core/tests/test_ndarray.py b/asdf/_tests/tags/core/tests/test_ndarray.py index 53b63da3d..6b36c5535 100644 --- a/asdf/_tests/tags/core/tests/test_ndarray.py +++ b/asdf/_tests/tags/core/tests/test_ndarray.py @@ -21,7 +21,9 @@ def ndarray_tag(): af = asdf.AsdfFile() cvt = af.extension_manager.get_converter_for_type(np.ndarray) - full_tag = cvt.select_tag(np.zeros(0), af) + full_tag = cvt.select_tag(np.zeros(0), af._create_serialization_context()) + assert full_tag is not None + return full_tag.removeprefix("tag:stsci.edu:asdf/") diff --git a/asdf/_tests/test_config.py b/asdf/_tests/test_config.py index ea4f22832..1e297be6f 100644 --- a/asdf/_tests/test_config.py +++ b/asdf/_tests/test_config.py @@ -215,12 +215,17 @@ def test_resource_mappings(): def test_resource_manager(): + def assert_bytes_in(a: bytes, b: str | bytes) -> None: + """Helper function for checking if bytes `a` are in sequence `b` which may be string or bytes""" + assert isinstance(b, bytes) + assert a in b + with asdf.config_context() as config: # Initial resource manager should contain just the entry points resources: assert "http://stsci.edu/schemas/asdf/core/asdf-1.1.0" in config.resource_manager - assert ( - b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0" - in config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"] + assert_bytes_in( + b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0", + config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"], ) assert "http://somewhere.org/schemas/foo-1.0.0" not in config.resource_manager @@ -228,9 +233,9 @@ def test_resource_manager(): new_mapping = {"http://somewhere.org/schemas/foo-1.0.0": b"foo"} config.add_resource_mapping(new_mapping) assert "http://stsci.edu/schemas/asdf/core/asdf-1.1.0" in config.resource_manager - assert ( - b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0" - in config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"] + assert_bytes_in( + b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0", + config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"], ) assert "http://somewhere.org/schemas/foo-1.0.0" in config.resource_manager assert config.resource_manager["http://somewhere.org/schemas/foo-1.0.0"] == b"foo" @@ -238,9 +243,9 @@ def test_resource_manager(): # Remove a mapping and confirm that the manager no longer contains it: config.remove_resource_mapping(new_mapping) assert "http://stsci.edu/schemas/asdf/core/asdf-1.1.0" in config.resource_manager - assert ( - b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0" - in config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"] + assert_bytes_in( + b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0", + config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"], ) assert "http://somewhere.org/schemas/foo-1.0.0" not in config.resource_manager @@ -248,9 +253,9 @@ def test_resource_manager(): config.add_resource_mapping(new_mapping) config.reset_resources() assert "http://stsci.edu/schemas/asdf/core/asdf-1.1.0" in config.resource_manager - assert ( - b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0" - in config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"] + assert_bytes_in( + b"http://stsci.edu/schemas/asdf/core/asdf-1.1.0", + config.resource_manager["http://stsci.edu/schemas/asdf/core/asdf-1.1.0"], ) assert "http://somewhere.org/schemas/foo-1.0.0" not in config.resource_manager diff --git a/asdf/extension/_serialization_context.py b/asdf/extension/_serialization_context.py index 3dff3c209..9ab0160e7 100644 --- a/asdf/extension/_serialization_context.py +++ b/asdf/extension/_serialization_context.py @@ -9,9 +9,10 @@ if TYPE_CHECKING: from asdf import AsdfFile + from asdf._block.callback import DataCallback from asdf._block.manager import Manager as BlockManager from asdf.extension import ExtensionLike, ExtensionManager - from asdf.typing import ArrayStorage, BlockDataCallback, Compression, NDArray + from asdf.typing import ArrayStorage, BlockAttrCallback, BlockDataCallback, Compression, NDArray from asdf.versioning import AsdfVersion @@ -97,7 +98,7 @@ def _extensions_used(self) -> set[ExtensionProxy]: """ return self.__extensions_used - def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> BlockDataCallback: + def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> BlockAttrCallback: """ Generate a callable that when called will read data from an ASDF block at the provided index. @@ -318,7 +319,7 @@ def assign_blocks(self) -> None: # assigned object self.assign_object(None) - def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> BlockDataCallback: + def get_block_data_callback(self, index: int, key: BlockKey | None = None) -> DataCallback: if key is None: if self._cb is not None: # this operation has already accessed a block without using diff --git a/asdf/tags/core/ndarray.py b/asdf/tags/core/ndarray.py index ab1c5fe09..bcd9dc0a5 100644 --- a/asdf/tags/core/ndarray.py +++ b/asdf/tags/core/ndarray.py @@ -12,9 +12,11 @@ from asdf._jsonschema import ValidationError if typing.TYPE_CHECKING: + from collections.abc import Sequence + import numpy.typing as npt - from asdf.typing import NDArray + from asdf.typing import BlockAttrCallback, NDArray _STRUCTURED_DATATYPE_KEYS = {"name", "datatype", "byteorder", "shape"} @@ -244,12 +246,12 @@ def ascii_to_unicode(x): return ascii_to_unicode(tolist(array)) -def inline_array_relax_empty_shape(array: NDArray, shape: tuple[int | str, ...] | None) -> NDArray: +def inline_array_relax_empty_shape(array: NDArray, shape: Sequence[int | str] | None) -> NDArray: if shape is None or any(isinstance(s, str) for s in shape): return array # unfortunately above lines do not trigger correct type-narrowing - shape = typing.cast("tuple[int, ...]", shape) + shape = typing.cast("Sequence[int]", shape) if array.size == 0 and np.prod(shape) == 0: array = array.reshape(shape) @@ -257,7 +259,17 @@ def inline_array_relax_empty_shape(array: NDArray, shape: tuple[int | str, ...] class NDArrayType: - def __init__(self, source, shape, dtype, offset, strides, order, mask, data_callback=None): + def __init__( + self, + source: str | int | list[Any], + shape: Sequence[int | str] | None, + dtype, + offset, + strides, + order, + mask, + data_callback: BlockAttrCallback | None = None, + ): self._source = source self._data_callback = data_callback self._array = None @@ -292,7 +304,7 @@ def _make_array(self) -> npt.NDArray[Any]: # closed and replaced. We need to check here and re-generate # the array if necessary, otherwise we risk segfaults when # memory mapping. - if self._array is not None: + if self._array is not None and self._data_callback is not None: base = util.get_array_base(self._array) if isinstance(base, np.memmap) and isinstance(base.base, mmap.mmap): # check if the underlying mmap matches the one generated by generic_io @@ -309,14 +321,18 @@ def _make_array(self) -> npt.NDArray[Any]: if self._array is not None: return self._array + # Currently arrays are always constructed with either inline data or a data callback + # This should ideally be enforced by the type system + data_callback = typing.cast("BlockAttrCallback", self._data_callback) + if isinstance(self._source, str): # we need to keep _source as a str to allow stdatamodels to # support AsdfInFits - data = self._data_callback() + data = data_callback() else: # cached data is used here so that multiple NDArrayTypes will all use # the same base array - data = self._data_callback(_attr="cached_data") + data = data_callback(_attr="cached_data") if hasattr(data, "base") and isinstance(data.base, mmap.mmap) and data.base.closed: msg = "ASDF file has already been closed. Can not get the data." @@ -408,9 +424,10 @@ def dtype(self): return self._make_array().dtype - def __len__(self): + def __len__(self) -> int: if self._array is None: - return self._shape[0] + # Array being empty means that its an inline array + return typing.cast("list[int]", self._shape)[0] return len(self._make_array()) diff --git a/asdf/typing.py b/asdf/typing.py index b648113ba..2dd548c3a 100644 --- a/asdf/typing.py +++ b/asdf/typing.py @@ -4,7 +4,7 @@ from collections.abc import Callable, Mapping, Sequence from pathlib import Path -from typing import Any, Literal, TypeAlias +from typing import Any, Literal, Protocol, TypeAlias, overload import numpy as np import numpy.typing as npt @@ -70,11 +70,27 @@ #: A 1-D byte numpy array used to read and write block data ByteArray1D: TypeAlias = np.ndarray[tuple[int], np.dtype[np.uint8]] -_Array = TypeVar("_Array", default=NDArray) +NotSetType = Literal[_NOT_SET_TYPE.NOT_SET] + +_Array_co = TypeVar("_Array_co", default=NDArray, covariant=True) + + +class ArrayCallback(Protocol[_Array_co]): + """A callback that returns a numpy array""" + + def __call__(self) -> _Array_co: ... + -#: A callback that returns a numpy array -ArrayCallback = Callable[[], _Array] #: A callback that returns a `ByteArray1D` BlockDataCallback = ArrayCallback[ByteArray1D] -NotSetType = Literal[_NOT_SET_TYPE.NOT_SET] + +class BlockAttrCallback(ArrayCallback[ByteArray1D], Protocol): + """A data callback that provides access to low-level block attributes.""" + + @overload + def __call__(self) -> ByteArray1D: ... + @overload + def __call__(self, _attr: str) -> Any: ... + + def __call__(self, _attr: str | None = None) -> ByteArray1D | Any: ... From d81ec717e28d930eac21f21b87deadfa7a738297 Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Wed, 12 Aug 2026 09:33:13 -0400 Subject: [PATCH 6/8] Fixed documentation --- asdf/extension/_compressor.py | 2 +- asdf/extension/_extension.py | 2 +- asdf/typing.py | 2 ++ docs/asdf/extending/compressors.rst | 17 +++++++++++------ docs/asdf/user_api/asdf_typing.rst | 2 +- docs/conf.py | 6 +++++- 6 files changed, 21 insertions(+), 10 deletions(-) diff --git a/asdf/extension/_compressor.py b/asdf/extension/_compressor.py index a5cb8050a..1042375af 100644 --- a/asdf/extension/_compressor.py +++ b/asdf/extension/_compressor.py @@ -99,7 +99,7 @@ class Compressor(Compress, Decompress): Abstract base class for plugins that compress binary data. Implementing classes must provide the ``labels`` property, and - at least one of the `compress()` and `decompress()` methods. + at least one of the ``compress()`` and ``decompress()`` methods. May also provide a constructor. """ diff --git a/asdf/extension/_extension.py b/asdf/extension/_extension.py index 126d5e72a..11fb0403c 100644 --- a/asdf/extension/_extension.py +++ b/asdf/extension/_extension.py @@ -43,7 +43,7 @@ class Extension(ExtensionLike): """ Abstract base class defining an extension to ASDF. - Implementing classes must provide the `extension_uri`. + Implementing classes must provide the ``extension_uri``. Other properties are optional. """ diff --git a/asdf/typing.py b/asdf/typing.py index 2dd548c3a..6d451aca0 100644 --- a/asdf/typing.py +++ b/asdf/typing.py @@ -16,8 +16,10 @@ from asdf.versioning import AsdfVersion __all__ = [ + "ArrayCallback", "ArrayStorage", "AsdfVersionLike", + "BlockAttrCallback", "BlockDataCallback", "ByteArray1D", "Compression", diff --git a/docs/asdf/extending/compressors.rst b/docs/asdf/extending/compressors.rst index a3cf38e44..ca986bdcf 100644 --- a/docs/asdf/extending/compressors.rst +++ b/docs/asdf/extending/compressors.rst @@ -20,15 +20,14 @@ a Compressor in an extension. The Compressor interface ======================== -Every Compressor implementation must provide one required property -and two required methods: +The Compressor interface is expressed through three related protocols: -`Compressor.label` - A 4-byte compression code. This code is used +`CompressionPlugin.label` - A 4-byte compression code. This code is used by users to select a compression algorithm and also stored in the binary block header to identify the algorithm that was applied to the block's data. -`Compressor.compress` - The method that transforms the block's bytes +`Compress.compress` - The method that transforms the block's bytes before they are written to an ASDF file. The positional argument is a `memoryview` object which is guaranteed to be 1D and contiguous. Compressors must be prepared to handle `memoryview.itemsize` > 1. @@ -37,7 +36,7 @@ to tune the compression algorithm. ``compress`` methods have no return value and instead are expected to yield bytes-like values until the input data has been fully compressed. -`Compressor.decompress` - The method that transforms the block's bytes +`Decompress.decompress` - The method that transforms the block's bytes after they are read from an ASDF file. The first positional argument is an `~collections.abc.Iterable` of bytes-like objects that each contain a chunk of the compressed input data. The second positional @@ -45,6 +44,12 @@ argument is a pre-allocated output array where the decompressed bytes should be written. The method is expected to return the number of bytes written to the output array. +Implementing `CompressionPlugin` is required, along with at least one of +`Compress` or `Decompress`. +As with all Python protocols, implementing the required methods is sufficient +to be considered a subclass. However, classes can also explicitly inherit from +`Compress` and/or `Decompress` to opt-in to function signature verification. + Entry point performance considerations ====================================== @@ -55,6 +60,6 @@ compressor module or ``__init__`` method that lingers will introduce a delay to the initial call to `asdf.open`. For that reason, we recommend that compressor authors minimize the number of imports that occur in the module containing the Compressor implementation, and defer imports of compression libraries to inside -the `Compressor.compress` and `Compressor.decompress` methods. This will +the `Compress.compress` and `Decompress.decompress` methods. This will prevent the library from ever being imported when reading ASDF files that do not utilize the Compressor's algorithm. diff --git a/docs/asdf/user_api/asdf_typing.rst b/docs/asdf/user_api/asdf_typing.rst index f889f434f..0fa342556 100644 --- a/docs/asdf/user_api/asdf_typing.rst +++ b/docs/asdf/user_api/asdf_typing.rst @@ -5,4 +5,4 @@ asdf.typing Module .. automodapi:: asdf.typing :include-all-objects: :no-inheritance-diagram: - :skip: NDArray, ByteArray1D, BlockDataCallback + :skip: NDArray, ByteArray1D, BlockDataCallback, ArrayCallback, BlockAttrCallback diff --git a/docs/conf.py b/docs/conf.py index fc9db30e4..c57a61dcc 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -54,8 +54,12 @@ nitpick_ignore = [ ("py:class", "yaml.representer.RepresenterError"), ("py:class", "yaml.error.YAMLError"), - # Ignore since its not part of the public API + # Ignore since they're not part of the public API ("py:attr", "asdf.util._NOT_SET_TYPE.NOT_SET"), + ("py:class", "BlockAttrCallback"), + ("py:class", "BlockManager"), + ("py:class", "BlockKey"), + ("py:class", "asdf._block.key.Key"), # Needed because sphinx breaks trying to process `asdf.typing.NDArray` for some reason ("py:class", "NDArray"), ("py:class", "ByteArray1D"), From 0372b366713cdb9a7a15451014d46519f5325738 Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Wed, 12 Aug 2026 10:53:52 -0400 Subject: [PATCH 7/8] Deprecated `Compressor` class --- asdf/_tests/test_compression.py | 8 ++++---- asdf/_tests/test_extension.py | 4 ++-- asdf/extension/_compressor.py | 3 +++ asdf/extension/_extension.py | 14 +++++++------- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/asdf/_tests/test_compression.py b/asdf/_tests/test_compression.py index 9ec7f10bb..4ccb31c20 100644 --- a/asdf/_tests/test_compression.py +++ b/asdf/_tests/test_compression.py @@ -10,7 +10,7 @@ import asdf from asdf import _compression, config_context from asdf._tests import _helpers as helpers -from asdf.extension import Compressor, Extension +from asdf.extension import Compress, Decompress, Extension if typing.TYPE_CHECKING: from asdf.generic_io import GenericFile @@ -208,15 +208,15 @@ def test_nonnative_endian_compression(tmp_path): _roundtrip(tmp_path, {"ledata": ledata, "bedata": bedata}, "lz4") -class LzmaCompressor(Compressor): +class LzmaCompressor(Compress, Decompress): def compress(self, data, **kwargs): comp = lzma.compress(data, **kwargs) yield comp - def decompress(self, blocks, out, **kwargs): + def decompress(self, data, out, **kwargs): decompressor = lzma.LZMADecompressor(**kwargs) i = 0 - for block in blocks: + for block in data: decomp = decompressor.decompress(block) out[i : i + len(decomp)] = decomp i += len(decomp) diff --git a/asdf/_tests/test_extension.py b/asdf/_tests/test_extension.py index 3d179d5b9..49da50bf3 100644 --- a/asdf/_tests/test_extension.py +++ b/asdf/_tests/test_extension.py @@ -11,7 +11,7 @@ from asdf import AsdfFile, config_context from asdf.exceptions import AsdfManifestURIMismatchWarning, AsdfSerializationError, ValidationError from asdf.extension import ( - Compressor, + Compress, Converter, ConverterProxy, Extension, @@ -112,7 +112,7 @@ def select_tag(self, obj, tags, ctx): return "select_tag result" -class MinimalCompressor(Compressor): +class MinimalCompressor(Compress): @staticmethod def compress(data): return b"" diff --git a/asdf/extension/_compressor.py b/asdf/extension/_compressor.py index 1042375af..57af8a32e 100644 --- a/asdf/extension/_compressor.py +++ b/asdf/extension/_compressor.py @@ -14,6 +14,8 @@ import abc from typing import TYPE_CHECKING, Protocol, runtime_checkable +from typing_extensions import deprecated + if TYPE_CHECKING: from collections.abc import Iterable, Iterator @@ -94,6 +96,7 @@ def decompress(self, data: Iterable[bytes], out: memoryview, **kwargs) -> int: raise NotImplementedError +@deprecated("Use Compress and Decompress protocols instead") class Compressor(Compress, Decompress): """ Abstract base class for plugins that compress binary data. diff --git a/asdf/extension/_extension.py b/asdf/extension/_extension.py index 11fb0403c..e6a8559eb 100644 --- a/asdf/extension/_extension.py +++ b/asdf/extension/_extension.py @@ -5,9 +5,9 @@ from packaging.specifiers import SpecifierSet +from asdf.extension import Compress, CompressionPlugin, Decompress from asdf.util import get_class_name -from ._compressor import Compressor from ._converter import ConverterProxy from ._tag import TagDefinition from ._validator import Validator @@ -104,14 +104,14 @@ def tags(self) -> Iterable[str | TagDefinition]: return [] @property - def compressors(self) -> Iterable[Compressor]: + def compressors(self) -> Iterable[CompressionPlugin]: """ - Get the `asdf.extension.Compressor` instances for + Get the `asdf.extension.CompressionPlugin` instances for compression schemes supported by this extension. Returns ------- - iterable of asdf.extension.Compressor + iterable of asdf.extension.CompressionPlugin """ return [] @@ -216,7 +216,7 @@ def __init__(self, delegate: ExtensionLike, package_name=None, package_version=N self._compressors = [] if hasattr(self._delegate, "compressors"): for compressor in self._delegate.compressors: - if not isinstance(compressor, Compressor): + if not isinstance(compressor, (Compress, Decompress)): msg = "Extension property 'compressors' must contain instances of asdf.extension.Compressor" raise TypeError(msg) self._compressors.append(compressor) @@ -279,13 +279,13 @@ def converters(self) -> list[ConverterProxy]: return self._converters @property - def compressors(self) -> list[Compressor]: + def compressors(self) -> list[CompressionPlugin]: """ Get the extension's compressors. Returns ------- - list of asdf.extension.Compressor + list of asdf.extension.CompressionPlugin """ return self._compressors From dce9804b2bce009a3a1f2d7881194107885eb466 Mon Sep 17 00:00:00 2001 From: Sydney Duckworth Date: Wed, 12 Aug 2026 11:02:17 -0400 Subject: [PATCH 8/8] Added changelog entries --- changes/2114.doc.rst | 1 + changes/2114.feature.1.rst | 1 + changes/2114.feature.rst | 1 + changes/2114.removal.rst | 1 + 4 files changed, 4 insertions(+) create mode 100644 changes/2114.doc.rst create mode 100644 changes/2114.feature.1.rst create mode 100644 changes/2114.feature.rst create mode 100644 changes/2114.removal.rst diff --git a/changes/2114.doc.rst b/changes/2114.doc.rst new file mode 100644 index 000000000..dcb3e0032 --- /dev/null +++ b/changes/2114.doc.rst @@ -0,0 +1 @@ +Updated documentation on writing compression plugins to reference new protocols. diff --git a/changes/2114.feature.1.rst b/changes/2114.feature.1.rst new file mode 100644 index 000000000..bedcccb3d --- /dev/null +++ b/changes/2114.feature.1.rst @@ -0,0 +1 @@ +Changed `asdf.extension.Converter` from an abstract class to a type-checkable protocol. diff --git a/changes/2114.feature.rst b/changes/2114.feature.rst new file mode 100644 index 000000000..f3850d8c7 --- /dev/null +++ b/changes/2114.feature.rst @@ -0,0 +1 @@ +Added new type-checkable protocols for compression plugins: `asdf.extension.CompressionPlugin`, `asdf.extension.Compress`, and `asdf.extension.Decompress`. diff --git a/changes/2114.removal.rst b/changes/2114.removal.rst new file mode 100644 index 000000000..3b7bcc74e --- /dev/null +++ b/changes/2114.removal.rst @@ -0,0 +1 @@ +Deprecated `asdf.extension.Compressor` class in favor of new compression protocols.