Merge pull request #2031 from sydduckworth/type-checking-fixes

Added type hints to `_asdf.py` and `config.py`

Author: sydduckworth

asdf/_asdf.py

@@ -74,4 +74,5 @@ def defragment(input_, output=None, resolve_references=False, compress=None):
         ff2 = AsdfFile(ff)
         if resolve_references:
             ff2.resolve_references()
-        ff2.write_to(output, all_array_storage="internal", all_array_compression=compress)
+        if output is not None:
+            ff2.write_to(output, all_array_storage="internal", all_array_compression=compress)

asdf/_commands/defragment.py

@@ -9,7 +9,7 @@
 from .exceptions import AsdfWarning
 
 
-def validate(compression):
+def validate(compression: str | bytes | None) -> str | None:
     """
     Validate the compression string.
 
@@ -318,7 +318,7 @@ def compress(fd, data, compression, config=None):
     # 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
-    data = memoryview(np.frombuffer(data, dtype=data.format))  # type: ignore
+    data = memoryview(np.frombuffer(data, dtype=data.format))  # pyrefly: ignore[bad-argument-type]
     if not data.contiguous:
         # the data will be contiguous by construction, but better safe than sorry!
         raise ValueError(data.contiguous)

asdf/_compression.py

@@ -1,7 +1,7 @@
 from io import BytesIO
 
 from asdf._asdf import AsdfFile, open_asdf
-from asdf.util import NotSet
+from asdf.util import NOT_SET
 
 __all__ = ["dump", "dumps", "load", "loads"]
 
@@ -12,9 +12,9 @@ def dump(
     *,
     version=None,
     extensions=None,
-    all_array_storage=NotSet,
-    all_array_compression=NotSet,
-    compression_kwargs=NotSet,
+    all_array_storage=NOT_SET,
+    all_array_compression=NOT_SET,
+    compression_kwargs=NOT_SET,
     pad_blocks=False,
     custom_schema=None,
     write_checksums=True,
@@ -75,9 +75,9 @@ def dumps(
     *,
     version=None,
     extensions=None,
-    all_array_storage=NotSet,
-    all_array_compression=NotSet,
-    compression_kwargs=NotSet,
+    all_array_storage=NOT_SET,
+    all_array_compression=NOT_SET,
+    compression_kwargs=NOT_SET,
     pad_blocks=False,
     custom_schema=None,
     write_checksums=True,

asdf/_dump.py

@@ -2,9 +2,14 @@
 Low-level input/output routines for AsdfFile instances
 """
 
+from __future__ import annotations
+
 import contextlib
 import io
 import pathlib
+from typing import TypeVar
+
+from asdf.versioning import AsdfVersion
 
 from . import _version, constants, generic_io, versioning, yamlutil
 from ._block.manager import Manager as BlockManager
@@ -27,7 +32,7 @@ def get_asdf_library_info():
     )
 
 
-def parse_header_line(line):
+def parse_header_line(line: str) -> AsdfVersion:
     """
     Parses the header line in a ASDF file to obtain the ASDF version.
     """
@@ -84,7 +89,11 @@ def read_comment_section(fd):
     return comments
 
 
-def find_asdf_version_in_comments(comments, default=None):
+# Required for type narrowing in `find_asdf_version_in_comments``
+MaybeVer = TypeVar("MaybeVer", AsdfVersion, None)
+
+
+def find_asdf_version_in_comments(comments: list[str], default: MaybeVer = None) -> MaybeVer:
     for comment in comments:
         parts = comment.split()
         if len(parts) == 2 and parts[0] == constants.ASDF_STANDARD_COMMENT:

asdf/_io.py

@@ -1,15 +1,10 @@
 from collections.abc import Mapping, MutableMapping, Sequence
 from urllib.parse import urlsplit
+from importlib import resources
+
 import itertools
 import json
 import re
-import sys
-
-# The files() API was added in Python 3.9.
-if sys.version_info >= (3, 9):  # pragma: no cover
-    from importlib import resources
-else:  # pragma: no cover
-    import importlib_resources as resources  # type: ignore
 
 
 class URIDict(MutableMapping):

asdf/_jsonschema/_utils.py

@@ -2,6 +2,7 @@
 
 import io
 import os
+from typing import Any
 
 import numpy as np
 import pytest
@@ -10,6 +11,7 @@
 from numpy.testing import assert_array_equal
 
 import asdf
+import asdf.versioning
 from asdf import constants, generic_io
 from asdf._block import io as bio
 from asdf.exceptions import AsdfBlockIndexWarning
@@ -57,6 +59,8 @@ def test_invalid_array_storage():
     tree = {"my_array": my_array}
     ff = asdf.AsdfFile(tree)
     with pytest.raises(ValueError, match=r"array_storage must be one of.*"):
+        # Intentionally incorrect storage type
+        # pyrefly: ignore[bad-argument-type]
         ff.set_array_storage(my_array, "foo")
 
 
@@ -883,7 +887,9 @@ def test_add_block_before_fully_loaded(tmp_path):
 @pytest.mark.parametrize("all_array_storage", ["internal", "external", "inline"])
 @pytest.mark.parametrize("all_array_compression", [None, "", "zlib", "bzp2", "lz4", "input"])
 @pytest.mark.parametrize("compression_kwargs", [None, {}])
-def test_write_to_update_storage_options(tmp_path, all_array_storage, all_array_compression, compression_kwargs):
+def test_write_to_update_storage_options(
+    tmp_path, all_array_storage, all_array_compression, compression_kwargs: dict[str, Any]
+):
     if all_array_compression == "bzp2" and compression_kwargs is not None:
         compression_kwargs = {"compresslevel": 1}
 

asdf/_tests/test_array_blocks.py

@@ -52,7 +52,7 @@ def url_mapping(self):
         return self._url_mapping
 
     @property
-    def extension_uri(self):
+    def extension_uri(self) -> str | None:
         return self._extension_uri
 
     @property
@@ -120,7 +120,8 @@ def test_asdf_file_extensions():
     msg = r"[The extensions parameter must be an extension.*, Extension must implement the Extension interface]"
     for arg in (object(), [object()]):
         with pytest.raises(TypeError, match=msg):
-            AsdfFile(extensions=arg)
+            # Ignore typing here because type is intentionally wrong
+            AsdfFile(extensions=arg)  # pyrefly: ignore[bad-argument-type]
 
 
 def test_asdf_file_version_requirement():
@@ -174,6 +175,8 @@ def test_open_asdf_extensions(tmp_path):
 
     msg = r"[The extensions parameter must be an extension.*, Extension must implement the Extension interface]"
     for arg in (object(), [object()]):
+        # Intentionally incorrect extension type
+        # pyrefly: ignore[bad-argument-type]
         with pytest.raises(TypeError, match=msg), open_asdf(path, extensions=arg) as af:
             pass
 

asdf/_tests/test_asdf.py

@@ -1,6 +1,7 @@
 import io
 import lzma
 import os
+from typing import Any
 
 import numpy as np
 import pytest
@@ -26,7 +27,13 @@ def _get_sparse_tree():
     return {"science_data": arr}
 
 
-def _roundtrip(tmp_path, tree, compression=None, write_options=None, read_options=None):
+def _roundtrip(
+    tmp_path,
+    tree,
+    compression=None,
+    write_options: dict[str, Any] | None = None,
+    read_options: dict[str, Any] | None = None,
+):
     read_options = {} if read_options is None else read_options
     write_options = {} if write_options is None else write_options.copy()
     write_options.update(all_array_compression=compression)
@@ -69,6 +76,8 @@ def test_invalid_compression():
     tree = _get_large_tree()
     ff = asdf.AsdfFile(tree)
     with pytest.raises(ValueError, match=r"Invalid compression type: foo"):
+        # Intentionally incorrect compression type
+        # pyrefly: ignore[bad-argument-type]
         ff.set_array_compression(tree["science_data"], "foo")
     with pytest.raises(ValueError, match=r"Unknown compression type: .*"):
         _compression._get_compressor("foo")

asdf/_tests/test_compression.py

@@ -117,6 +117,8 @@ def test_all_array_storage():
         config.all_array_storage = None
         assert get_config().all_array_storage is None
         with pytest.raises(ValueError, match=r"Invalid value for all_array_storage"):
+            # Intentionally incorrect argument type
+            # pyrefly: ignore[bad-argument-type]
             config.all_array_storage = "foo"
 
 
@@ -139,6 +141,8 @@ def test_all_array_compression_kwargs():
         config.all_array_compression_kwargs = None
         assert get_config().all_array_compression_kwargs is None
         with pytest.raises(ValueError, match=r"Invalid value for all_array_compression_kwargs"):
+            # Intentionally incorrect argument type
+            # pyrefly: ignore[bad-argument-type]
             config.all_array_compression_kwargs = "foo"