Additional type hinting

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,
@@ -329,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

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.

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

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)

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)
 

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)

asdf/_core/_converters/external_reference.py

@@ -1,19 +1,32 @@
+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,
             "datatype": obj.dtype,
             "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"]))