diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 52700d2a3..2abd34bb2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -85,6 +85,16 @@ jobs: - libopenjp2-7 - graphviz + type_check: + uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@v3 # zizmor: ignore[unpinned-uses] + with: + submodules: false + pytest: false + toxdeps: tox-pypi-filter + envs: | + - linux: mypy + - linux: pyright + cron: if: | github.event_name == 'workflow_dispatch' || ( @@ -111,7 +121,7 @@ jobs: github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'Run publish') ) - needs: [test, docs] + needs: [test, docs, type_check] uses: OpenAstronomy/github-actions-workflows/.github/workflows/publish_pure_python.yml@v2 # zizmor: ignore[unpinned-uses] with: python-version: '3.13' diff --git a/changelog/950.feature.rst b/changelog/950.feature.rst new file mode 100644 index 000000000..910e269c8 --- /dev/null +++ b/changelog/950.feature.rst @@ -0,0 +1 @@ +Add type and type checking to NDCube diff --git a/docs/explaining_ndcube/slicing.rst b/docs/explaining_ndcube/slicing.rst index 779d57b20..3d70a787b 100644 --- a/docs/explaining_ndcube/slicing.rst +++ b/docs/explaining_ndcube/slicing.rst @@ -267,7 +267,7 @@ If we want our region of interest to only apply to a single sub-cube, and we ind >>> single_cube_roi.array_axis_physical_types [('custom:pos.helioprojective.lat', 'custom:pos.helioprojective.lon'), ('em.wl',)] -However, as with numpy slicing, we can induce the slicing operation to return an `~ndcube.NDCubeSequence` by supplying a length-1 `slice` to the sequence axis, rather than an `int`. +However, as with numpy slicing, we can induce the slicing operation to return an `~ndcube.NDCubeSequence` by supplying a length-1 ``slice`` to the sequence axis, rather than an `int`. This sequence will still represent the same region of interest from the same single sub-cube, but the sequence axis will have a length of 1, rather than be removed. .. code-block:: python @@ -309,7 +309,7 @@ This can be achieved by entering: [('custom:pos.helioprojective.lat', 'custom:pos.helioprojective.lon'), ('em.wl',)] This returns the same `~ndcube.NDCube` as above. -However, also as above, we can induce the return type to be an `~ndcube.NDCubeSequence` by supplying a length-1 `slice`. +However, also as above, we can induce the return type to be an `~ndcube.NDCubeSequence` by supplying a length-1 ``slice``. As before, the same region of interest from the same sub-cube is represented, just with sequence and common axes of length 1. .. code-block:: python diff --git a/docs/nitpick-exceptions b/docs/nitpick-exceptions index cf555dca3..03cd03c57 100644 --- a/docs/nitpick-exceptions +++ b/docs/nitpick-exceptions @@ -41,6 +41,7 @@ py:class an object providing a view on D's values py:class D[k] if k in D, else d. d defaults to None. py:class None. Remove all items from D. py:class None. Update D from mapping/iterable E and F. +py:class None. Update D from dict/iterable E and F. py:class v, remove specified key and return the corresponding value. # This comes from Map.wcs @@ -77,6 +78,14 @@ py:class numpy.core.records.recarray py:class numpy.ma.core.MaskedArray py:class numpy.ma.mvoid py:class numpy.void + +# numpy.typing.NDArray's real __module__ is this private path, not the +# public numpy.typing name, so autodoc can't resolve it via intersphinx +py:class numpy._typing._array_like.NDArray + +# ndcube.utils.cube.WCSType is an internal typing.TypeAlias, not a +# documented public class +py:class WCSType py:class pandas.DataFrame py:class xmlrpc.client.Error py:class xmlrpc.client.Fault diff --git a/ndcube/asdf/converters/compoundwcs_converter.py b/ndcube/asdf/converters/compoundwcs_converter.py index 1f733fb4a..6a4429a07 100644 --- a/ndcube/asdf/converters/compoundwcs_converter.py +++ b/ndcube/asdf/converters/compoundwcs_converter.py @@ -1,17 +1,19 @@ +from typing import Any + from asdf.extension import Converter -class CompoundConverter(Converter): +class CompoundConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/compoundwcs-*"] types = ["ndcube.wcs.wrappers.compound_wcs.CompoundLowLevelWCS"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.wcs.wrappers import CompoundLowLevelWCS - return CompoundLowLevelWCS(*node["wcs"], mapping=node.get("mapping"), pixel_atol=node.get("atol")) + return CompoundLowLevelWCS(*node["wcs"], mapping=node.get("mapping"), pixel_atol=node.get("atol")) # type: ignore[arg-type] - def to_yaml_tree(self, compoundwcs, tag, ctx): - node = {} + def to_yaml_tree(self, compoundwcs: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["wcs"] = compoundwcs._wcs node["mapping"] = compoundwcs.mapping.mapping node["atol"] = compoundwcs.atol diff --git a/ndcube/asdf/converters/extracoords_converter.py b/ndcube/asdf/converters/extracoords_converter.py index 0a0c74b69..9c8000d2c 100644 --- a/ndcube/asdf/converters/extracoords_converter.py +++ b/ndcube/asdf/converters/extracoords_converter.py @@ -1,29 +1,31 @@ +from typing import Any + from asdf.extension import Converter -class ExtraCoordsConverter(Converter): +class ExtraCoordsConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/extra_coords/extra_coords/extracoords-*"] types = ["ndcube.extra_coords.extra_coords.ExtraCoords"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.extra_coords.extra_coords import ExtraCoords extra_coords = ExtraCoords() - extra_coords._wcs = node.get("wcs") - extra_coords._mapping = node.get("mapping") - extra_coords._lookup_tables = node.get("lookup_tables", []) - extra_coords._dropped_tables = node.get("dropped_tables") - extra_coords._ndcube = node.get("ndcube") + extra_coords._wcs = node.get("wcs") # pyright: ignore[reportPrivateUsage] + extra_coords._mapping = node.get("mapping") # pyright: ignore[reportPrivateUsage] + extra_coords._lookup_tables = node.get("lookup_tables", []) # pyright: ignore[reportPrivateUsage] + extra_coords._dropped_tables = node.get("dropped_tables", []) # pyright: ignore[reportPrivateUsage] + extra_coords._ndcube = node.get("ndcube") # pyright: ignore[reportPrivateUsage] return extra_coords - def to_yaml_tree(self, extracoords, tag, ctx): - node = {} - if extracoords._wcs is not None: - node["wcs"] = extracoords._wcs - if extracoords._mapping is not None: - node["mapping"] = extracoords._mapping - if extracoords._lookup_tables: - node["lookup_tables"] = extracoords._lookup_tables - if extracoords._dropped_tables is not None: - node["dropped_tables"] = extracoords._dropped_tables - node["ndcube"] = extracoords._ndcube + def to_yaml_tree(self, extracoords: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} + if extracoords._wcs is not None: # pyright: ignore[reportPrivateUsage] + node["wcs"] = extracoords._wcs # pyright: ignore[reportPrivateUsage] + if extracoords._mapping is not None: # pyright: ignore[reportPrivateUsage] + node["mapping"] = extracoords._mapping # pyright: ignore[reportPrivateUsage] + if extracoords._lookup_tables: # pyright: ignore[reportPrivateUsage] + node["lookup_tables"] = extracoords._lookup_tables # pyright: ignore[reportPrivateUsage] + if extracoords._dropped_tables is not None: # pyright: ignore[reportPrivateUsage] + node["dropped_tables"] = extracoords._dropped_tables # pyright: ignore[reportPrivateUsage] + node["ndcube"] = extracoords._ndcube # pyright: ignore[reportPrivateUsage] return node diff --git a/ndcube/asdf/converters/globalcoords_converter.py b/ndcube/asdf/converters/globalcoords_converter.py index 9700a5890..5f0e717da 100644 --- a/ndcube/asdf/converters/globalcoords_converter.py +++ b/ndcube/asdf/converters/globalcoords_converter.py @@ -1,24 +1,26 @@ +from typing import Any + from asdf.extension import Converter -class GlobalCoordsConverter(Converter): +class GlobalCoordsConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/global_coords/globalcoords-*"] types = ["ndcube.global_coords.GlobalCoords"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.global_coords import GlobalCoords globalcoords = GlobalCoords() if "internal_coords" in node: - globalcoords._internal_coords = node["internal_coords"] - globalcoords._ndcube = node["ndcube"] + globalcoords._internal_coords = node["internal_coords"] # pyright: ignore[reportPrivateUsage] + globalcoords._ndcube = node["ndcube"] # pyright: ignore[reportPrivateUsage] return globalcoords - def to_yaml_tree(self, globalcoords, tag, ctx): - node = {} - node["ndcube"] = globalcoords._ndcube - if globalcoords._internal_coords: - node["internal_coords"] = globalcoords._internal_coords + def to_yaml_tree(self, globalcoords: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} + node["ndcube"] = globalcoords._ndcube # pyright: ignore[reportPrivateUsage] + if globalcoords._internal_coords: # pyright: ignore[reportPrivateUsage] + node["internal_coords"] = globalcoords._internal_coords # pyright: ignore[reportPrivateUsage] return node diff --git a/ndcube/asdf/converters/ndcollection_converter.py b/ndcube/asdf/converters/ndcollection_converter.py index 1bbe44a95..67d7522e6 100644 --- a/ndcube/asdf/converters/ndcollection_converter.py +++ b/ndcube/asdf/converters/ndcollection_converter.py @@ -1,19 +1,21 @@ +from typing import Any + from asdf.extension import Converter -class NDCollectionConverter(Converter): +class NDCollectionConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/ndcollection-*"] types = ["ndcube.ndcollection.NDCollection"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.ndcollection import NDCollection - aligned_axes = list(node.get("aligned_axes", {}).values()) - aligned_axes = tuple(tuple(lst) for lst in aligned_axes) + aligned_axes_list = list(node.get("aligned_axes", {}).values()) + aligned_axes = tuple(tuple(lst) for lst in aligned_axes_list) return NDCollection(node["items"], meta=node.get("meta"), aligned_axes=aligned_axes) - def to_yaml_tree(self, ndcollection, tag, ctx): - node = {} + def to_yaml_tree(self, ndcollection: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["items"] = dict(ndcollection) if ndcollection.meta is not None: node["meta"] = ndcollection.meta diff --git a/ndcube/asdf/converters/ndcube_converter.py b/ndcube/asdf/converters/ndcube_converter.py index fa7a03b45..7b81d9ee6 100644 --- a/ndcube/asdf/converters/ndcube_converter.py +++ b/ndcube/asdf/converters/ndcube_converter.py @@ -1,13 +1,14 @@ import warnings +from typing import Any from asdf.extension import Converter -class NDCubeConverter(Converter): +class NDCubeConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/ndcube-*"] types = ["ndcube.ndcube.NDCube"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.ndcube import NDCube ndcube = NDCube( @@ -19,13 +20,13 @@ def from_yaml_tree(self, node, tag, ctx): uncertainty=node.get("uncertainty"), ) if "extra_coords" in node: - ndcube._extra_coords = node["extra_coords"] + ndcube._extra_coords = node["extra_coords"] # pyright: ignore[reportPrivateUsage] if "global_coords" in node: - ndcube._global_coords = node["global_coords"] + ndcube._global_coords = node["global_coords"] # pyright: ignore[reportPrivateUsage] return ndcube - def to_yaml_tree(self, ndcube, tag, ctx): + def to_yaml_tree(self, ndcube: Any, tag: str, ctx: Any) -> dict[str, Any]: """ Notes ----- @@ -41,7 +42,7 @@ def to_yaml_tree(self, ndcube, tag, ctx): This ensures that users are aware of potentially important information that is not included in the serialized output. """ - node = {} + node: dict[str, Any] = {} node["data"] = ndcube.data # NDData always has .wcs as a high level wcs node["wcs"] = ndcube.wcs.low_level_wcs diff --git a/ndcube/asdf/converters/ndcubesequence_converter.py b/ndcube/asdf/converters/ndcubesequence_converter.py index 64896dbaf..e533b689c 100644 --- a/ndcube/asdf/converters/ndcubesequence_converter.py +++ b/ndcube/asdf/converters/ndcubesequence_converter.py @@ -1,19 +1,21 @@ +from typing import Any + from asdf.extension import Converter -class NDCubeSequenceConverter(Converter): +class NDCubeSequenceConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/ndcubesequence-*"] types = ["ndcube.ndcube_sequence.NDCubeSequence"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.ndcube_sequence import NDCubeSequence return NDCubeSequence(node["data"], meta=node.get("meta"), common_axis=node.get("common_axis")) - def to_yaml_tree(self, ndcseq, tag, ctx): - node = {} + def to_yaml_tree(self, ndcseq: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["data"] = ndcseq.data if ndcseq.meta is not None: node["meta"] = ndcseq.meta diff --git a/ndcube/asdf/converters/ndmeta_converter.py b/ndcube/asdf/converters/ndmeta_converter.py index f8fa011b3..ae3bbfa6c 100644 --- a/ndcube/asdf/converters/ndmeta_converter.py +++ b/ndcube/asdf/converters/ndmeta_converter.py @@ -1,24 +1,26 @@ +from typing import Any + import numpy as np from asdf.extension import Converter -class NDMetaConverter(Converter): +class NDMetaConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/meta/ndmeta-*"] types = ["ndcube.meta.NDMeta"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.meta import NDMeta axes = {k: np.array(v) for k, v in node["axes"].items()} meta = NDMeta(node["meta"], node["key_comments"], axes, node["data_shape"]) - meta._original_meta = node["original_meta"] + meta._original_meta = node["original_meta"] # pyright: ignore[reportPrivateUsage] return meta - def to_yaml_tree(self, meta, tag, ctx): - node = {} + def to_yaml_tree(self, meta: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["meta"] = dict(meta) node["key_comments"] = meta.key_comments node["axes"] = meta.axes node["data_shape"] = meta.data_shape - node["original_meta"] = meta._original_meta # not the MappingProxy object + node["original_meta"] = meta._original_meta # not the MappingProxy object # pyright: ignore[reportPrivateUsage] return node diff --git a/ndcube/asdf/converters/reorderedwcs_converter.py b/ndcube/asdf/converters/reorderedwcs_converter.py index 81982f8f9..8521f9b8f 100644 --- a/ndcube/asdf/converters/reorderedwcs_converter.py +++ b/ndcube/asdf/converters/reorderedwcs_converter.py @@ -1,11 +1,13 @@ +from typing import Any + from asdf.extension import Converter -class ReorderedConverter(Converter): +class ReorderedConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/reorderedwcs-*"] types = ["ndcube.wcs.wrappers.reordered_wcs.ReorderedLowLevelWCS"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.wcs.wrappers import ReorderedLowLevelWCS return ReorderedLowLevelWCS( @@ -14,8 +16,8 @@ def from_yaml_tree(self, node, tag, ctx): world_order=node["world_order"], ) - def to_yaml_tree(self, reorderedwcs, tag, ctx): - node = {} + def to_yaml_tree(self, reorderedwcs: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["wcs"] = reorderedwcs._wcs node["pixel_order"] = reorderedwcs._pixel_order node["world_order"] = reorderedwcs._world_order diff --git a/ndcube/asdf/converters/resampled_converter.py b/ndcube/asdf/converters/resampled_converter.py index 465c5f7bd..67431e585 100644 --- a/ndcube/asdf/converters/resampled_converter.py +++ b/ndcube/asdf/converters/resampled_converter.py @@ -1,11 +1,13 @@ +from typing import Any + from asdf.extension import Converter -class ResampledConverter(Converter): +class ResampledConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/resampledwcs-*"] types = ["ndcube.wcs.wrappers.resampled_wcs.ResampledLowLevelWCS"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.wcs.wrappers import ResampledLowLevelWCS return ResampledLowLevelWCS( @@ -14,8 +16,8 @@ def from_yaml_tree(self, node, tag, ctx): factor=node["factor"], ) - def to_yaml_tree(self, resampledwcs, tag, ctx): - node = {} + def to_yaml_tree(self, resampledwcs: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["wcs"] = resampledwcs._wcs node["factor"] = resampledwcs._factor node["offset"] = resampledwcs._offset diff --git a/ndcube/asdf/converters/tablecoord_converter.py b/ndcube/asdf/converters/tablecoord_converter.py index 31c2dba4a..5c0e57cf5 100644 --- a/ndcube/asdf/converters/tablecoord_converter.py +++ b/ndcube/asdf/converters/tablecoord_converter.py @@ -1,11 +1,13 @@ +from typing import Any + from asdf.extension import Converter -class TimeTableCoordConverter(Converter): +class TimeTableCoordConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/timetablecoordinate-*"] types = ["ndcube.extra_coords.table_coord.TimeTableCoordinate"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.extra_coords.table_coord import TimeTableCoordinate names = node.get("names") @@ -18,8 +20,8 @@ def from_yaml_tree(self, node, tag, ctx): reference_time=reference_time, ) - def to_yaml_tree(self, timetablecoordinate, tag, ctx): - node = {} + def to_yaml_tree(self, timetablecoordinate: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["table"] = timetablecoordinate.table if timetablecoordinate.names: node["names"] = timetablecoordinate.names @@ -30,11 +32,11 @@ def to_yaml_tree(self, timetablecoordinate, tag, ctx): return node -class QuantityTableCoordinateConverter(Converter): +class QuantityTableCoordinateConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/quantitytablecoordinate-*"] types = ["ndcube.extra_coords.table_coord.QuantityTableCoordinate"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.extra_coords.table_coord import QuantityTableCoordinate names = node.get("names") @@ -42,11 +44,11 @@ def from_yaml_tree(self, node, tag, ctx): physical_types = node.get("physical_types") quantitytablecoordinate = QuantityTableCoordinate(*node["table"], names=names, physical_types=physical_types) quantitytablecoordinate.unit = node["unit"] - quantitytablecoordinate.mesh = mesh + quantitytablecoordinate.mesh = mesh # type: ignore[assignment] return quantitytablecoordinate - def to_yaml_tree(self, quantitytablecoordinate, tag, ctx): - node = {} + def to_yaml_tree(self, quantitytablecoordinate: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["unit"] = quantitytablecoordinate.unit node["table"] = quantitytablecoordinate.table if quantitytablecoordinate.names: @@ -58,20 +60,20 @@ def to_yaml_tree(self, quantitytablecoordinate, tag, ctx): return node -class SkyCoordTableCoordinateConverter(Converter): +class SkyCoordTableCoordinateConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/skycoordtablecoordinate-*"] types = ["ndcube.extra_coords.table_coord.SkyCoordTableCoordinate"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.extra_coords.table_coord import SkyCoordTableCoordinate names = node.get("names") mesh = node.get("mesh") physical_types = node.get("physical_types") - return SkyCoordTableCoordinate(node["table"], mesh=mesh, names=names, physical_types=physical_types) + return SkyCoordTableCoordinate(node["table"], mesh=mesh, names=names, physical_types=physical_types) # type: ignore[arg-type] - def to_yaml_tree(self, skycoordinatetablecoordinate, tag, ctx): - node = {} + def to_yaml_tree(self, skycoordinatetablecoordinate: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} node["table"] = skycoordinatetablecoordinate.table if skycoordinatetablecoordinate.names: node["names"] = skycoordinatetablecoordinate.names @@ -82,19 +84,19 @@ def to_yaml_tree(self, skycoordinatetablecoordinate, tag, ctx): return node -class MultipleTableCoordinateConverter(Converter): +class MultipleTableCoordinateConverter(Converter): # type: ignore[misc] tags = ["tag:sunpy.org:ndcube/extra_coords/table_coord/multipletablecoordinate-*"] types = ["ndcube.extra_coords.table_coord.MultipleTableCoordinate"] - def from_yaml_tree(self, node, tag, ctx): + def from_yaml_tree(self, node: dict[str, Any], tag: str, ctx: Any) -> Any: from ndcube.extra_coords.table_coord import MultipleTableCoordinate mtc = MultipleTableCoordinate(*node["table_coords"]) - mtc._dropped_coords = node["dropped_coords"] + mtc._dropped_coords = node["dropped_coords"] # pyright: ignore[reportPrivateUsage] return mtc - def to_yaml_tree(self, multipletablecoordinate, tag, ctx): - node = {} - node["table_coords"] = multipletablecoordinate._table_coords + def to_yaml_tree(self, multipletablecoordinate: Any, tag: str, ctx: Any) -> dict[str, Any]: + node: dict[str, Any] = {} + node["table_coords"] = multipletablecoordinate._table_coords # pyright: ignore[reportPrivateUsage] node["dropped_coords"] = multipletablecoordinate._dropped_coords return node diff --git a/ndcube/asdf/entry_points.py b/ndcube/asdf/entry_points.py index ade044f41..4d7ce0de9 100644 --- a/ndcube/asdf/entry_points.py +++ b/ndcube/asdf/entry_points.py @@ -2,12 +2,14 @@ This file contains the entry points for asdf. """ import importlib.resources as importlib_resources +from typing import Any +from collections.abc import Sequence from asdf.extension import ManifestExtension from asdf.resource import DirectoryResourceMapping -def get_resource_mappings(): +def get_resource_mappings() -> Sequence[Any]: """ Get the resource mapping instances for myschemas and manifests. This method is registered with the @@ -17,7 +19,7 @@ def get_resource_mappings(): ------- list of collections.abc.Mapping """ - from ndcube.asdf import resources + from ndcube.asdf import resources # pyright: ignore[reportAttributeAccessIssue] resources_root = importlib_resources.files(resources) return [ DirectoryResourceMapping( @@ -27,7 +29,7 @@ def get_resource_mappings(): ] -def get_extensions(): +def get_extensions() -> list[Any]: """ Get the list of extensions. """ diff --git a/ndcube/extra_coords/extra_coords.py b/ndcube/extra_coords/extra_coords.py index 85243597f..e3ceaed3f 100644 --- a/ndcube/extra_coords/extra_coords.py +++ b/ndcube/extra_coords/extra_coords.py @@ -1,8 +1,8 @@ import abc -from typing import Any +from typing import TYPE_CHECKING, Any from numbers import Integral from functools import reduce, partial -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import numpy as np @@ -10,7 +10,7 @@ from astropy.coordinates import SkyCoord from astropy.time import Time from astropy.wcs import WCS -from astropy.wcs.wcsapi import BaseHighLevelWCS +from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS from astropy.wcs.wcsapi.high_level_wcs_wrapper import HighLevelWCSWrapper from astropy.wcs.wcsapi.wrappers.sliced_wcs import SlicedLowLevelWCS, sanitize_slices @@ -19,11 +19,15 @@ from .table_coord import ( BaseTableCoordinate, MultipleTableCoordinate, + NamesType, QuantityTableCoordinate, SkyCoordTableCoordinate, TimeTableCoordinate, ) +if TYPE_CHECKING: + from ndcube.ndcube import NDCubeABC + __all__ = ['ExtraCoords', 'ExtraCoordsABC'] @@ -47,11 +51,11 @@ class ExtraCoordsABC(abc.ABC): """ @abc.abstractmethod def add(self, - name: str | Iterable[str], + name: NamesType, array_dimension: int | Iterable[int], lookup_table: Any, - physical_types: str | Iterable[str] = None, - **kwargs): + physical_types: NamesType = None, + **kwargs: Any) -> None: """ Add a coordinate to this `~ndcube.ExtraCoords` based on a lookup table. @@ -72,24 +76,25 @@ def add(self, """ @abc.abstractmethod - def keys(self) -> Iterable[str]: + def keys(self) -> Iterable[str] | None: """ The world axis names for all the coordinates in the extra coords. """ @property @abc.abstractmethod - def mapping(self) -> Iterable[tuple[int, int]]: + def mapping(self) -> Iterable[int]: """ The mapping between the array dimensions and pixel dimensions. - This is an iterable of ``(array_dimension, pixel_dimension)`` pairs - of length equal to the number of pixel dimensions in the extra coords. + This is an iterable of pixel dimension indices, one per array dimension + in array-axis order, of length equal to the number of pixel dimensions + in the extra coords. """ @property @abc.abstractmethod - def wcs(self) -> BaseHighLevelWCS: + def wcs(self) -> BaseHighLevelWCS | None: """ A WCS object representing the world coordinates described by this ``ExtraCoords``. @@ -105,9 +110,16 @@ def wcs(self) -> BaseHighLevelWCS: @property @abc.abstractmethod - def is_empty(self): + def is_empty(self) -> bool: """Return True if no extra coords present, else return False.""" + @property + @abc.abstractmethod + def dropped_world_dimensions(self) -> dict[str, Any]: + """ + Return an APE-14 like representation of any sliced out world dimensions. + """ + @abc.abstractmethod def __getitem__(self, item: str | int | slice | Iterable[str | int | slice]) -> "ExtraCoordsABC": """ @@ -120,6 +132,15 @@ def __getitem__(self, item: str | int | slice | Iterable[str | int | slice]) -> not supported. """ + @abc.abstractmethod + def resample(self, factor: float | Iterable[float], offset: float | Iterable[float] = 0, + ndcube: Any = None, **kwargs: Any) -> "ExtraCoordsABC": + """ + Resample all extra coords by given factors in array-index-space. + + One resample factor must be supplied for each array axis in array-axis order. + """ + class ExtraCoords(ExtraCoordsABC): """ @@ -140,23 +161,25 @@ class ExtraCoords(ExtraCoordsABC): """ - def __init__(self, ndcube=None): + def __init__(self, ndcube: "NDCubeABC | None" = None) -> None: super().__init__() # Setup private attributes - self._wcs = None - self._mapping = None + self._wcs: BaseLowLevelWCS | BaseHighLevelWCS | None = None + self._mapping: Iterable[int] | None = None # Lookup tables is a list of (pixel_dim, LookupTableCoord) to allow for # one pixel dimension having more than one lookup coord. - self._lookup_tables = [] - self._dropped_tables = [] + self._lookup_tables: list[tuple[Any, BaseTableCoordinate]] = [] + self._dropped_tables: list[BaseTableCoordinate] = [] # We need a reference to the parent NDCube self._ndcube = ndcube @classmethod - def from_lookup_tables(cls, names, pixel_dimensions, lookup_tables, physical_types=None): + def from_lookup_tables(cls, names: Sequence[str], pixel_dimensions: Sequence[Any], + lookup_tables: Sequence[Any], + physical_types: Sequence[Any] | None = None) -> "ExtraCoords": """ Construct a new ExtraCoords instance from lookup tables. @@ -204,7 +227,8 @@ def from_lookup_tables(cls, names, pixel_dimensions, lookup_tables, physical_typ return extra_coords - def add(self, name, array_dimension, lookup_table, physical_types=None, **kwargs): + def add(self, name: NamesType, array_dimension: int | Iterable[int], lookup_table: Any, + physical_types: NamesType = None, **kwargs: Any) -> None: # docstring in ABC if self._wcs is not None: @@ -234,13 +258,13 @@ def add(self, name, array_dimension, lookup_table, physical_types=None, **kwargs key=lambda x: x[0] if isinstance(x[0], Integral) else x[0][0]) @property - def _name_lut_map(self): + def _name_lut_map(self) -> dict[Any, tuple[Any, BaseTableCoordinate]]: """ Map of world names to the corresponding `.LookupTableCoord` """ return {lut[1].wcs.world_axis_names: lut for lut in self._lookup_tables} - def keys(self): + def keys(self) -> tuple[str, ...] | None: # docstring in ABC if not self.wcs: return () @@ -248,10 +272,10 @@ def keys(self): return tuple(self.wcs.world_axis_names) if self.wcs.world_axis_names else None @property - def mapping(self): + def mapping(self) -> tuple[int, ...]: # docstring in ABC if self._mapping: - return self._mapping + return tuple(self._mapping) # If mapping is not set but lookup_tables is empty then the extra # coords is empty, so there is no mapping. @@ -261,12 +285,12 @@ def mapping(self): # The mapping is from the array index (position in the list) to the # pixel dimensions (numbers in the list) lts = [list([lt[0]] if isinstance(lt[0], Integral) else lt[0]) for lt in self._lookup_tables] - converter = partial(convert_between_array_and_pixel_axes, naxes=len(self._ndcube.shape)) + converter = partial(convert_between_array_and_pixel_axes, naxes=len(self._ndcube.shape)) # type: ignore[union-attr] pixel_indicies = [list(converter(np.array(ids))) for ids in lts] - return tuple(reduce(list.__add__, pixel_indicies)) + return tuple(reduce(list.__add__, pixel_indicies)) # type: ignore[arg-type] @mapping.setter - def mapping(self, mapping): + def mapping(self, mapping: Iterable[int]) -> None: if self._mapping is not None: raise AttributeError("Can't set mapping if a mapping has already been specified.") @@ -284,7 +308,7 @@ def mapping(self, mapping): self._mapping = mapping @property - def wcs(self): + def wcs(self) -> BaseHighLevelWCS | None: # docstring in ABC if self._wcs is not None: return self._wcs @@ -292,14 +316,13 @@ def wcs(self): if not self._lookup_tables: return None - tcoords = {lt[1] for lt in self._lookup_tables} # created a sorted list of unique items - _tmp = set() # a temporary set - tcoords = [x[1] for x in self._lookup_tables if x[1] not in _tmp and _tmp.add(x[1]) is None] + _tmp: set[Any] = set() # a temporary set + tcoords = [x[1] for x in self._lookup_tables if x[1] not in _tmp and _tmp.add(x[1]) is None] # type: ignore[func-returns-value] return MultipleTableCoordinate(*tcoords).wcs @wcs.setter - def wcs(self, wcs): + def wcs(self, wcs: BaseLowLevelWCS | BaseHighLevelWCS) -> None: if self._wcs is not None: raise AttributeError( "Can't set wcs if a WCS has already been specified." @@ -319,13 +342,13 @@ def wcs(self, wcs): self._wcs = wcs @property - def is_empty(self): + def is_empty(self) -> bool: # docstring in ABC if not self._wcs and not self._lookup_tables: return True return False - def _getitem_string(self, item): + def _getitem_string(self, item: str) -> "ExtraCoords": """ Slice the Extracoords based on axis names. """ @@ -338,16 +361,16 @@ def _getitem_string(self, item): raise KeyError(f"Can't find the world axis named {item} in this ExtraCoords object.") - def _getitem_lookup_tables(self, item): + def _getitem_lookup_tables(self, item: Any) -> "ExtraCoords": """ Apply an array slice to the lookup tables. Returns a new ExtraCoords object with modified lookup tables. """ - dropped_tables = set() - new_lookup_tables = set() - ndims = max([lut[0] if isinstance(lut[0], Integral) else max(lut[0]) - for lut in self._lookup_tables]) + 1 + dropped_tables: set[BaseTableCoordinate] = set() + new_lookup_tables: set[tuple[Any, BaseTableCoordinate]] = set() + ndims = int(max(lut[0] if isinstance(lut[0], Integral) else max(lut[0]) + for lut in self._lookup_tables)) + 1 # Determine how many dimensions will be dropped by slicing below each dimension. if isinstance(item, Integral): n_dropped_dims = np.ones(ndims, dtype=int) @@ -362,7 +385,7 @@ def _getitem_lookup_tables(self, item): lut_axes = (lut_axis,) if not isinstance(lut_axis, tuple) else lut_axis new_lut_axes = tuple(ax - n_dropped_dims[ax] for ax in lut_axes) lut_slice = tuple(item[i] for i in lut_axes) - if isinstance(lut_slice, tuple) and len(lut_slice) == 1: + if len(lut_slice) == 1: lut_slice = lut_slice[0] sliced_lut = lut[lut_slice] @@ -376,7 +399,10 @@ def _getitem_lookup_tables(self, item): new_extra_coords._dropped_tables = list(dropped_tables) return new_extra_coords - def _getitem_wcs(self, item): + def _getitem_wcs(self, item: Any) -> "ExtraCoords": + # Only called from __getitem__ when self._wcs is truthy, so self.wcs is not None here. + if self.wcs is None: + raise RuntimeError("wcs must not be None here.") item = sanitize_slices(item, self.wcs.pixel_n_dim) # It's valid to slice down the EC such that there is nothing left, @@ -393,7 +419,7 @@ def _getitem_wcs(self, item): new_ec.mapping = new_mapping return new_ec - def __getitem__(self, item): + def __getitem__(self, item: str | int | slice | Iterable[str | int | slice]) -> "ExtraCoords": # docstring in ABC if isinstance(item, str): return self._getitem_string(item) @@ -409,24 +435,25 @@ def __getitem__(self, item): return self @property - def dropped_world_dimensions(self): + def dropped_world_dimensions(self) -> dict[str, Any]: """ Return an APE-14 like representation of any sliced out world dimensions. """ if self._wcs: if isinstance(self._wcs, SlicedLowLevelWCS): - return self._wcs.dropped_world_dimensions + return self._wcs.dropped_world_dimensions # type: ignore[no-any-return] if self._lookup_tables or self._dropped_tables: mtc = MultipleTableCoordinate(*[lt[1] for lt in self._lookup_tables]) - mtc._dropped_coords = self._dropped_tables + mtc._dropped_coords = self._dropped_tables # pyright: ignore[reportPrivateUsage] return mtc.dropped_world_dimensions return {} - def resample(self, factor, offset=0, ndcube=None, **kwargs): + def resample(self, factor: float | Iterable[float], offset: float | Iterable[float] = 0, + ndcube: "NDCubeABC | None" = None, **kwargs: Any) -> "ExtraCoords": """ Resample all extra coords by given factors in array-index-space. @@ -445,8 +472,9 @@ def resample(self, factor, offset=0, ndcube=None, **kwargs): for each dimension. If a scalar, the grid will be shifted by the same amount in all dimensions. - ndcube: `~ndcube.NDCube` + ndcube: `~ndcube.NDCube`, optional The NDCube instance with which the output ExtraCoords object is associated. + Defaults to `None`. kwargs All remaining kwargs are passed to `numpy.interp`. @@ -469,41 +497,41 @@ def resample(self, factor, offset=0, ndcube=None, **kwargs): "Resampling a lookup-table-based ExtraCoords not yet implemented. " "Please raise an issue at https://github.com/sunpy/ndcube/issues " "if you need this functionality") - if np.isscalar(factor): - factor = [factor] * ndim - if len(factor) != ndim: + factor_seq: Sequence[float] = [factor] * ndim if np.isscalar(factor) else list(factor) # type: ignore[list-item, arg-type] + if len(factor_seq) != ndim: raise ValueError( "factor must be scalar or an iterable with length equal to number of cube " - f"dimensions: len(factor) = {len(factor)}; No. cube dimensions = {ndim}.") - if np.isscalar(offset): - offset = [offset] * ndim - if len(offset) != ndim: + f"dimensions: len(factor) = {len(factor_seq)}; No. cube dimensions = {ndim}.") + offset_seq: Sequence[float] = [offset] * ndim if np.isscalar(offset) else list(offset) # type: ignore[list-item, arg-type] + if len(offset_seq) != ndim: raise ValueError( "offset must be scalar or an iterable with length equal to number of cube " - f"dimensions: len(offset) = {len(offset)}; No. cube dimensions = {ndim}.") + f"dimensions: len(offset) = {len(offset_seq)}; No. cube dimensions = {ndim}.") # If ExtraCoords object built on WCS, resample using WCS insfrastructure if self._wcs is not None: new_ec.wcs = HighLevelWCSWrapper(ResampledLowLevelWCS(self._wcs.low_level_wcs, - factor, offset)) + factor_seq, offset_seq)) return new_ec # Else interpolate the lookup table coordinates. - factor = np.asarray(factor) + # self._wcs is None here (handled above), so the if/elif above guarantees + # self._ndcube was not None and cube_shape was set. + factor_arr = np.asarray(factor_seq) new_grids = [] - for c, d, f in zip(offset, cube_shape, factor): + for c, d, f in zip(offset_seq, cube_shape, factor_arr): # pyright: ignore[reportPossiblyUnboundVariable] x = np.arange(c, d+f, f) x = x[x <= d-1] new_grids.append(x) new_grids = np.array(new_grids, dtype=object) for array_axes, coord in self._lookup_tables: if np.isscalar(array_axes): - new_coord = coord.interpolate(new_grids[array_axes], **kwargs) + new_coord = coord.interpolate(new_grids[array_axes], **kwargs) # type: ignore[index] else: new_coord = coord.interpolate(*new_grids[np.asarray(array_axes)], **kwargs) - new_ec.add(coord.names, array_axes, new_coord, physical_types=coord.physical_types) + new_ec.add(coord.names, array_axes, new_coord, physical_types=coord.physical_types) # pyright: ignore[reportArgumentType] return new_ec @property - def cube_wcs(self): + def cube_wcs(self) -> CompoundLowLevelWCS: """Produce a WCS that describes the associated NDCube with just the extra coords. For NDCube pixel axes without any extra coord, dummy axes are inserted. @@ -521,16 +549,16 @@ def cube_wcs(self): dummy_wcs.wcs.cunit = ["pix"] * n_dummy_axes wcses.append(dummy_wcs) mapping += list(dummy_axes) - return CompoundLowLevelWCS(*wcses, mapping=mapping) + return CompoundLowLevelWCS(*wcses, mapping=tuple(mapping)) @property - def _cube_array_axes_without_extra_coords(self): + def _cube_array_axes_without_extra_coords(self) -> set[int]: """Return the array axes not associated with any extra coord.""" - return set(range(len(self._ndcube.shape))) - set(self.mapping) + return set(range(len(self._ndcube.shape))) - set(self.mapping) # type: ignore[union-attr] - def __str__(self): + def __str__(self) -> str: classname = self.__class__.__name__ - elements = [f"{', '.join(table.names)} ({axes}) {table.physical_types}: {table}" + elements = [f"{', '.join(table.names or ())} ({axes}) {table.physical_types}: {table}" for axes, table in self._lookup_tables] length = len(classname) + 2 * len(elements) + sum(len(e) for e in elements) if length > np.get_printoptions()['linewidth']: @@ -540,5 +568,5 @@ def __str__(self): return f"{classname}({joiner.join(elements)})" - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self}" diff --git a/ndcube/extra_coords/table_coord.py b/ndcube/extra_coords/table_coord.py index e975393d4..5e9d00eeb 100644 --- a/ndcube/extra_coords/table_coord.py +++ b/ndcube/extra_coords/table_coord.py @@ -1,17 +1,20 @@ import abc import copy import uuid +from typing import Any from numbers import Integral from functools import partial from collections import defaultdict +from collections.abc import Iterable, Sequence import gwcs import gwcs.coordinate_frames as cf import numpy as np +import numpy.typing as npt import astropy.units as u from astropy.coordinates import SkyCoord -from astropy.modeling import models +from astropy.modeling import Model, models from astropy.modeling.models import tabular_model from astropy.modeling.tabular import Tabular1D, Tabular2D, _Tabular from astropy.time import Time @@ -22,6 +25,10 @@ except ImportError: pass +# Names/physical types are usually given as a single `str` but can be a +# `list` of `str` when a coordinate has multiple components (e.g. SkyCoord). +NamesType = str | list[str] | None + __all__ = [ "BaseTableCoordinate", "MultipleTableCoordinate", @@ -31,7 +38,7 @@ ] -class Length1Tabular(_Tabular): +class Length1Tabular(_Tabular): # type: ignore[misc] _input_units_allow_dimensionless = True _has_inverse_bounding_box = True _separable = True @@ -39,11 +46,13 @@ class Length1Tabular(_Tabular): n_inputs = 1 n_outputs = 1 - lookup_table = np.zeros([1]) - points = np.zeros([1]) + lookup_table: u.Quantity = np.zeros([1]) + points: u.Quantity = np.zeros([1]) - def __init__(self, points=None, lookup_table=None, point_width=None, value_width=None, - method='linear', bounds_error=True, fill_value=np.nan, **kwargs): + def __init__(self, points: u.Quantity | None = None, lookup_table: u.Quantity | None = None, + point_width: u.Quantity | None = None, value_width: u.Quantity | None = None, + method: str = 'linear', bounds_error: bool = True, fill_value: u.Quantity | float = np.nan, + **kwargs: Any) -> None: """Create a Length-1 1-D Tabular model. Parameters @@ -60,18 +69,16 @@ def __init__(self, points=None, lookup_table=None, point_width=None, value_width Other parameters are defined by the parent class. """ - if len(lookup_table) != 1: + if lookup_table is None or len(lookup_table) != 1: raise ValueError("lookup_table must have length 1.") super().__init__(points=points, lookup_table=lookup_table, method=method, bounds_error=bounds_error, fill_value=fill_value, **kwargs) - self._value_width = value_width # Width of point in world units. - if self._value_width is None: - self._value_width = 0 * self.lookup_table.unit - self._point_width = point_width # Width of point in point units. - if self._point_width is None: - self._point_width = 1 * self.points[0].unit - - def evaluate(self, x): + # Width of point in world units. + self._value_width: u.Quantity = value_width if value_width is not None else 0 * self.lookup_table.unit + # Width of point in point units. + self._point_width: u.Quantity = point_width if point_width is not None else 1 * self.points[0].unit + + def evaluate(self, x: u.Quantity) -> u.Quantity: output = np.full(x.shape, self.fill_value) diff = abs(x - self.points[0]) margin = self._point_width / 2 @@ -83,7 +90,7 @@ def evaluate(self, x): return output * self.lookup_table.unit @property - def inverse(self): + def inverse(self) -> "InverseLength1Tabular": return InverseLength1Tabular(points=self.points[0], lookup_table=self.lookup_table, point_width=self._point_width, value_width=self._value_width, method=self.method, bounds_error=self.bounds_error, @@ -96,7 +103,7 @@ class InverseLength1Tabular(Length1Tabular): This is the opposite direction to Length1Tabular. """ - def __init__(self, **kwargs): + def __init__(self, **kwargs: Any) -> None: # Same inputs as Length1Tabular points = kwargs.pop("points", None) lookup_table = kwargs.pop("lookup_table", None) @@ -105,13 +112,14 @@ def __init__(self, **kwargs): super().__init__(points=lookup_table, lookup_table=points, point_width=value_width, value_width=point_width, **kwargs) - def evaluate(self, x): + def evaluate(self, x: u.Quantity) -> u.Quantity: # When calling evaluate with a bounding box, astropy strips the units. x = u.Quantity(x, unit=self.input_units['x'], copy=False) return super().evaluate(x) -def _generate_generic_frame(naxes, unit, names=None, physical_types=None): +def _generate_generic_frame(naxes: int, unit: u.UnitBase | Iterable[u.UnitBase], + names: NamesType = None, physical_types: NamesType = None) -> cf.CoordinateFrame: """ Generate a simple frame, where all axes have the same type and unit. """ @@ -140,7 +148,8 @@ def _generate_generic_frame(naxes, unit, names=None, physical_types=None): axes_names=names, name=name, axis_physical_types=physical_types) -def _generate_tabular(lookup_table, interpolation='linear', points_unit=u.pix, **kwargs): +def _generate_tabular(lookup_table: u.Quantity, interpolation: str = 'linear', + points_unit: u.UnitBase = u.pix, **kwargs: Any) -> Model: """ Generate a Tabular model class and instance. """ @@ -182,7 +191,7 @@ def _generate_tabular(lookup_table, interpolation='linear', points_unit=u.pix, * return t -def _generate_compound_model(*lookup_tables, mesh=True): +def _generate_compound_model(*lookup_tables: u.Quantity, mesh: bool = True) -> Model: """ Takes a set of quantities and returns a ND compound model. """ @@ -198,7 +207,7 @@ def _generate_compound_model(*lookup_tables, mesh=True): return models.Mapping(mapping) | model -def _model_from_quantity(lookup_tables, mesh=False): +def _model_from_quantity(lookup_tables: Sequence[u.Quantity], mesh: bool = False) -> Model: if len(lookup_tables) > 1: return _generate_compound_model(*lookup_tables, mesh=mesh) @@ -218,19 +227,22 @@ class BaseTableCoordinate(abc.ABC): coordinates, meaning it can have multiple gWCS frames. """ - def __init__(self, *tables, mesh=False, names=None, physical_types=None): - self.table = tables + def __init__(self, *tables: Any, mesh: bool = False, names: NamesType = None, + physical_types: NamesType = None) -> None: + self.table: Any = tables self.mesh = mesh - self.names = names if not isinstance(names, str) else [names] - self.physical_types = physical_types if not isinstance(physical_types, str) else [physical_types] - self._dropped_world_dimensions = defaultdict(list) + self.names: list[str] | None = names if not isinstance(names, str) else [names] + self.physical_types: list[str] | None = ( + physical_types if not isinstance(physical_types, str) else [physical_types] + ) + self._dropped_world_dimensions: defaultdict[str, Any] = defaultdict(list) self._dropped_world_dimensions["world_axis_object_classes"] = {} @abc.abstractmethod - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "BaseTableCoordinate": pass # pragma: no cover - def __and__(self, other): + def __and__(self, other: Any) -> "BaseTableCoordinate": if not isinstance(other, BaseTableCoordinate): return NotImplemented @@ -242,25 +254,25 @@ def __and__(self, other): return MultipleTableCoordinate(self, other) - def __str__(self): + def __str__(self) -> str: header = f"{self.__class__.__name__} {self.names or ''} {self.physical_types or '[None]'}:" content = str(self.table).lstrip('(').rstrip(',)') if len(header) + len(content) >= np.get_printoptions()['linewidth']: return f'{header}\n{content}' return f'{header} {content}' - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self}" @property @abc.abstractmethod - def n_inputs(self): + def n_inputs(self) -> int: """ Number of pixel dimensions in this table. """ @abc.abstractmethod - def is_scalar(self): + def is_scalar(self) -> bool: """ Return a boolean if this coordinate is a scalar. @@ -270,20 +282,20 @@ def is_scalar(self): @property @abc.abstractmethod - def frame(self): + def frame(self) -> cf.CoordinateFrame: """ Generate the Frame for this LookupTable. """ @property @abc.abstractmethod - def model(self): + def model(self) -> Model: """ Generate the Astropy Model for this LookupTable. """ @property - def wcs(self): + def wcs(self) -> gwcs.WCS: """ A gWCS object representing all the coordinates. """ @@ -293,9 +305,15 @@ def wcs(self): output_frame=self.frame) @property - def dropped_world_dimensions(self): + def dropped_world_dimensions(self) -> dict[str, Any]: return self._dropped_world_dimensions + @abc.abstractmethod + def interpolate(self, *new_array_grids: npt.NDArray[Any], **kwargs: Any) -> "BaseTableCoordinate": + """ + Interpolate this coordinate to new array index grids. + """ + class QuantityTableCoordinate(BaseTableCoordinate): """ @@ -315,14 +333,14 @@ class QuantityTableCoordinate(BaseTableCoordinate): Custom names for the components of the QuantityTableCoord. If provided, a name must be given for each input Quantity. - physical_types: str` or `list` of `str` + physical_types: `str` or `list` of `str` Physical types for the components of the QuantityTableCoord. If provided, a physical type must be given for each input Quantity. Physical types of the components of the SkyCoord. If provided, a physical type must be given for each component. """ - def __init__(self, *tables, names=None, physical_types=None): + def __init__(self, *tables: u.Quantity, names: NamesType = None, physical_types: NamesType = None) -> None: if not all(isinstance(t, u.Quantity) for t in tables): raise TypeError("All tables must be astropy Quantity objects") if not all(t.unit.is_equivalent(tables[0].unit) for t in tables): @@ -347,7 +365,8 @@ def __init__(self, *tables, names=None, physical_types=None): super().__init__(*tables, mesh=True, names=names, physical_types=physical_types) - def _slice_table(self, i, table, item, new_components, whole_slice): + def _slice_table(self, i: int, table: u.Quantity, item: Any, + new_components: "defaultdict[str, Any]", whole_slice: Any) -> None: """ Apply a slice, or part of a slice to one of the quantity arrays. @@ -379,13 +398,13 @@ def _slice_table(self, i, table, item, new_components, whole_slice): if self.physical_types: new_components["physical_types"].append(self.physical_types[i]) - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "QuantityTableCoordinate": if isinstance(item, (slice, Integral)): item = (item,) if not (len(item) == len(self.table) or len(item) == self.table[0].ndim): raise ValueError("Can not slice with incorrect length") - new_components = defaultdict(list) + new_components: defaultdict[str, Any] = defaultdict(list) new_components["dropped_world_dimensions"] = copy.deepcopy(self._dropped_world_dimensions) for i, (ele, table) in enumerate(zip(item, self.table)): @@ -399,28 +418,28 @@ def __getitem__(self, item): return ret_table @property - def n_inputs(self): + def n_inputs(self) -> int: return len(self.table) - def is_scalar(self): + def is_scalar(self) -> bool: return all(t.shape == () for t in self.table) @property - def frame(self): + def frame(self) -> cf.CoordinateFrame: """ Generate the Frame for this LookupTable. """ return _generate_generic_frame(len(self.table), self.unit, self.names, self.physical_types) @property - def model(self): + def model(self) -> Model: """ Generate the Astropy Model for this LookupTable. """ return _model_from_quantity(self.table, True) @property - def ndim(self): + def ndim(self) -> int: """ Number of array dimensions to which this TableCoordinate corresponds. @@ -430,7 +449,7 @@ def ndim(self): return len(self.table) @property - def shape(self): + def shape(self) -> tuple[int, ...]: """ Shape of the array grid to which this TableCoordinate corresponds. @@ -439,7 +458,7 @@ def shape(self): """ return tuple(len(t) for t in self.table) - def interpolate(self, *new_array_grids, **kwargs): + def interpolate(self, *new_array_grids: npt.NDArray[Any], **kwargs: Any) -> "QuantityTableCoordinate": """ Interpolate QuantityTableCoordinate to new array index grids. @@ -499,7 +518,7 @@ class SkyCoordTableCoordinate(BaseTableCoordinate): Custom names for the components of the SkyCoord. If provided, a name must be given for each component. - physical_types: str` or `list` of `str` + physical_types: `str` or `list` of `str` Physical types of the components of the SkyCoord. If provided, a physical type must be given for each component. @@ -510,7 +529,8 @@ class SkyCoordTableCoordinate(BaseTableCoordinate): same length. """ - def __init__(self, *tables, mesh=False, names=None, physical_types=None): + def __init__(self, *tables: SkyCoord, mesh: bool = False, names: NamesType = None, + physical_types: NamesType = None) -> None: if not len(tables) == 1 and isinstance(tables[0], SkyCoord): raise ValueError("SkyCoordLookupTable can only be constructed from a single SkyCoord object") if mesh and tables[0].ndim > 1: @@ -533,14 +553,14 @@ def __init__(self, *tables, mesh=False, names=None, physical_types=None): self._slice = sanitize_slices(np.s_[...], self.n_inputs) @property - def n_inputs(self): + def n_inputs(self) -> int: return len(self.table.data.components) - def is_scalar(self): - return self.table.shape == () + def is_scalar(self) -> bool: + return bool(self.table.shape == ()) @staticmethod - def combine_slices(slice1, slice2): + def combine_slices(slice1: Any, slice2: Any) -> Any: ints = [isinstance(s, Integral) for s in (slice1, slice2)] if all(ints): raise ValueError("Can not combine two integers") @@ -548,7 +568,7 @@ def combine_slices(slice1, slice2): return (slice1, slice2)[ints.index(True)] return combine_slices(slice1, slice2) - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "SkyCoordTableCoordinate": # override the error for consistency try: sane_item = sanitize_slices(item, self.n_inputs) @@ -571,7 +591,7 @@ def __getitem__(self, item): return self @property - def frame(self): + def frame(self) -> cf.CoordinateFrame: """ Generate the Frame for this LookupTable. """ @@ -588,19 +608,19 @@ def frame(self): name="CelestialFrame") @property - def _sliced_components(self): + def _sliced_components(self) -> tuple[Any, ...]: return tuple(getattr(self.table.data, comp)[slc] for comp, slc in zip(self.table.data.components, self._slice)) @property - def model(self): + def model(self) -> Model: """ Generate the Astropy Model for this LookupTable. """ return _model_from_quantity(self._sliced_components, mesh=self.mesh) @property - def ndim(self): + def ndim(self) -> int: """ Number of array dimensions to which this TableCoordinate corresponds. @@ -610,10 +630,10 @@ def ndim(self): """ if self.mesh: return len(self.table.data.components) - return self.table.ndim + return int(self.table.ndim) @property - def shape(self): + def shape(self) -> tuple[int, ...]: """ Shape of the array grid to which this TableCoordinate corresponds. @@ -624,9 +644,10 @@ def shape(self): """ if self.mesh: return tuple(list(self.table.shape) * self.ndim) - return self.table.shape + return tuple(self.table.shape) - def interpolate(self, *new_array_grids, mesh_output=None, **kwargs): + def interpolate(self, *new_array_grids: npt.NDArray[Any], mesh_output: bool | None = None, + **kwargs: Any) -> "SkyCoordTableCoordinate": """ Interpolate SkyCoordTableCoordinate to new array index grids. @@ -680,11 +701,11 @@ def interpolate(self, *new_array_grids, mesh_output=None, **kwargs): for new_grid, old_grid, comp in zip(new_array_grids, old_array_grids, self._sliced_components)] elif ndim == 1: - new_components = [np.interp(*new_array_grids, *old_array_grids, comp, **kwargs) + new_components = [np.interp(*new_array_grids, *old_array_grids, comp, **kwargs) # type: ignore[call-overload] for comp in self._sliced_components] else: new_components = [ - scipy.interpolate.interpn(old_array_grids, component, new_array_grids, **kwargs) + scipy.interpolate.interpn(old_array_grids, component, new_array_grids, **kwargs) # pyright: ignore[reportPossiblyUnboundVariable] for component in self._sliced_components] # Build new SkyCoord and return new TableCoordinate based on it. @@ -710,7 +731,7 @@ class TimeTableCoordinate(BaseTableCoordinate): Custom names for the components of the SkyCoord. If provided, a name must be given for each component. - physical_types: str` or `list` of `str` + physical_types: `str` or `list` of `str` Physical types of the components of the SkyCoord. If provided, a physical type must be given for each component. @@ -719,7 +740,8 @@ class TimeTableCoordinate(BaseTableCoordinate): Default is first time coordinate in table input. """ - def __init__(self, *tables, names=None, physical_types=None, reference_time=None): + def __init__(self, *tables: Time, names: NamesType = None, physical_types: NamesType = None, + reference_time: Time | None = None) -> None: if not len(tables) == 1 and isinstance(tables[0], Time): raise ValueError("TimeLookupTable can only be constructed from a single Time object.") @@ -737,7 +759,7 @@ def __init__(self, *tables, names=None, physical_types=None, reference_time=None self.table = self.table[0] self.reference_time = reference_time or self.table[0] - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "TimeTableCoordinate": if not (isinstance(item, (slice, Integral)) or len(item) == 1): raise ValueError("Can not slice with incorrect length") @@ -747,14 +769,14 @@ def __getitem__(self, item): reference_time=self.reference_time) @property - def n_inputs(self): + def n_inputs(self) -> int: return 1 # The time table has to be one dimensional - def is_scalar(self): - return self.table.shape == () + def is_scalar(self) -> bool: + return bool(self.table.shape == ()) @property - def frame(self): + def frame(self) -> cf.CoordinateFrame: """ Generate the Frame for this LookupTable. """ @@ -764,7 +786,7 @@ def frame(self): name="TemporalFrame") @property - def model(self): + def model(self) -> Model: """ Generate the Astropy Model for this LookupTable. """ @@ -773,7 +795,9 @@ def model(self): return _model_from_quantity((deltas,), mesh=False) - def interpolate(self, new_array_grids, **kwargs): + # Deliberately takes a single grid (not *args like the base class) since a + # Time coordinate is always 1-D. + def interpolate(self, new_array_grids: npt.NDArray[Any], **kwargs: Any) -> "TimeTableCoordinate": # type: ignore[override] """ Interpolate TimeTableCoordinate to new array index grids. @@ -782,9 +806,10 @@ def interpolate(self, new_array_grids, **kwargs): Parameters ---------- new_array_grids: array-like - The array index values at which the the new values of the - coords are desired. A grid must be supplied for each pixel - axis (in array-axis order). All grids must be the same shape. + The array index values at which the new values of the coords are + desired. Since a TimeTableCoordinate is always 1-D, a single + array grid is given (not one per axis like the other + TableCoordinate classes' ``interpolate`` methods). Returns ------- @@ -826,15 +851,17 @@ class MultipleTableCoordinate(BaseTableCoordinate): operator. """ - def __init__(self, *table_coordinates): - if not all(isinstance(lt, BaseTableCoordinate) and + def __init__(self, *table_coordinates: BaseTableCoordinate) -> None: + # Runtime guard against callers not respecting the type hints, so isinstance + # isn't required as is statically implied by the annotation above + if not all(isinstance(lt, BaseTableCoordinate) and # pyright: ignore[reportUnnecessaryIsInstance] not (isinstance(lt, MultipleTableCoordinate)) for lt in table_coordinates): raise TypeError("All arguments must be BaseTableCoordinate instances, such as QuantityTableCoordinate, " "and not instances of MultipleTableCoordinate.") - self._table_coords = list(table_coordinates) - self._dropped_coords = [] + self._table_coords: list[BaseTableCoordinate] = list(table_coordinates) + self._dropped_coords: list[BaseTableCoordinate] = [] - def __str__(self): + def __str__(self) -> str: classname = self.__class__.__name__ length = len(classname) + sum(len(str(t)) for t in self._table_coords) + 10 if length > np.get_printoptions()['linewidth']: @@ -844,7 +871,7 @@ def __str__(self): return f"{classname}(tables=[{joiner.join([str(t) for t in self._table_coords])}])" - def __and__(self, other): + def __and__(self, other: Any) -> "MultipleTableCoordinate": if not isinstance(other, BaseTableCoordinate): return NotImplemented @@ -855,14 +882,14 @@ def __and__(self, other): return type(self)(*(self._table_coords + others)) - def __rand__(self, other): + def __rand__(self, other: Any) -> "MultipleTableCoordinate": # This method should never be called if the left hand operand is a MultipleTableCoordinate if not isinstance(other, BaseTableCoordinate) or isinstance(other, MultipleTableCoordinate): return NotImplemented return type(self)(*([other, *self._table_coords])) - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "MultipleTableCoordinate": if isinstance(item, (slice, Integral)): item = (item,) @@ -887,14 +914,14 @@ def __getitem__(self, item): return new @property - def n_inputs(self): + def n_inputs(self) -> int: return sum(t.n_inputs for t in self._table_coords) - def is_scalar(self): + def is_scalar(self) -> bool: return False @property - def model(self): + def model(self) -> Model: """ The combined astropy model for all the lookup tables. """ @@ -904,7 +931,7 @@ def model(self): return model @property - def frame(self): + def frame(self) -> cf.CoordinateFrame: """ The gWCS coordinate frame for all the lookup tables. """ @@ -924,7 +951,7 @@ def frame(self): return cf.CompositeFrame(frames) @staticmethod - def _from_high_level_coordinates(dropped_frame, *highlevel_coords): + def _from_high_level_coordinates(dropped_frame: Any, *highlevel_coords: Any) -> Any: """ This is a backwards compatibility wrapper for the new from_high_level_coordinates method in gwcs. @@ -935,8 +962,8 @@ def _from_high_level_coordinates(dropped_frame, *highlevel_coords): return quantities @property - def dropped_world_dimensions(self): - dropped_world_dimensions = defaultdict(list) + def dropped_world_dimensions(self) -> dict[str, Any]: + dropped_world_dimensions: defaultdict[str, Any] = defaultdict(list) dropped_world_dimensions["world_axis_object_classes"] = {} # Combine the dicts on the tables with our dict @@ -983,7 +1010,9 @@ def dropped_world_dimensions(self): return dropped_world_dimensions - def interpolate(self, new_array_grids, **kwargs): + # Deliberately takes a single grid (not *args like the base class); it is + # forwarded unchanged to each sub-coordinate's own interpolate(). + def interpolate(self, new_array_grids: npt.NDArray[Any], **kwargs: Any) -> "MultipleTableCoordinate": # type: ignore[override] """ Interpolate MultipleTableCoordinate to new array index grids. @@ -992,9 +1021,9 @@ def interpolate(self, new_array_grids, **kwargs): Parameters ---------- new_array_grids: array-like - The array index values at which the the new values of the - coords are desired. A grid must be supplied for each pixel - axis (in array-axis order). All grids must be the same shape. + A single array grid, forwarded unchanged to every sub-coordinate's + own ``interpolate`` method (not one grid per axis like + `~ndcube.extra_coords.table_coord.BaseTableCoordinate.interpolate`). Returns ------- @@ -1003,7 +1032,7 @@ def interpolate(self, new_array_grids, **kwargs): """ new_table_coordinates = [coord.interpolate(new_array_grids, **kwargs) - for coord in self.table_coords] + for coord in self._table_coords] new_obj = type(self)(*new_table_coordinates) new_obj._dropped_coords = self._dropped_coords return new_obj diff --git a/ndcube/extra_coords/tests/test_extra_coords.py b/ndcube/extra_coords/tests/test_extra_coords.py index e48190ace..7ca1b20bf 100644 --- a/ndcube/extra_coords/tests/test_extra_coords.py +++ b/ndcube/extra_coords/tests/test_extra_coords.py @@ -560,3 +560,21 @@ def test_length1_extra_coord(wave_lut): sec = ec[item] assert (sec.wcs.pixel_to_world(0) == wave_lut[item]).all() assert (sec.wcs.world_to_pixel(wave_lut[item])[0] == [0]).all() + + +def test_str_with_unnamed_table_coordinate(wave_lut): + # Regression test: ExtraCoords.__str__ used to do ', '.join(table.names), + # which crashes with TypeError if `names` is None. This happens whenever a + # BaseTableCoordinate instance is passed directly to `.add()`, since `add` + # only sets names via its own `name` argument when it constructs the + # coordinate itself, not when an already-built coordinate is passed in. + from ndcube.extra_coords.table_coord import QuantityTableCoordinate # NOQA + + ec = ExtraCoords() + coord = QuantityTableCoordinate(wave_lut) + assert coord.names is None + ec.add("wave", 0, coord) + + # Should not raise. + str(ec) + repr(ec) diff --git a/ndcube/global_coords.py b/ndcube/global_coords.py index 8c3c9f215..683cbbc7f 100644 --- a/ndcube/global_coords.py +++ b/ndcube/global_coords.py @@ -1,8 +1,8 @@ import abc import copy -from typing import Any +from typing import TYPE_CHECKING, Any from collections import OrderedDict, defaultdict -from collections.abc import Mapping +from collections.abc import Mapping, Iterator import numpy as np @@ -12,8 +12,18 @@ from ndcube.utils.wcs import validate_physical_types +if TYPE_CHECKING: + from ndcube.ndcube import NDCubeABC -class GlobalCoordsABC(Mapping): +# Coordinate names/physical types are usually `str`, but can be grouped into +# a `tuple` of `str` when several dropped world dimensions share one coordinate. +CoordKeyType = str | tuple[str, ...] +# A physical type is like CoordKeyType, but may also be None for coordinates +# added via `GlobalCoords.add` without a known UCD1+ physical type. +PhysicalTypeValue = CoordKeyType | None + + +class GlobalCoordsABC(Mapping[CoordKeyType, Any]): """ A structured representation of coordinate information applicable to a whole `~ndcube.ndcube.NDCubeABC`. @@ -30,7 +40,7 @@ class GlobalCoordsABC(Mapping): coordinates explicitly added will be shown. """ @abc.abstractmethod - def add(self, name: str, physical_type: str, coord: Any): + def add(self, name: str, physical_type: str | None, coord: Any) -> None: """ Add a new coordinate to the collection. @@ -49,32 +59,32 @@ def add(self, name: str, physical_type: str, coord: Any): """ @abc.abstractmethod - def remove(self, name: str): + def remove(self, name: str) -> None: """ Remove a coordinate from the collection. """ @property @abc.abstractmethod - def physical_types(self): + def physical_types(self) -> Mapping[CoordKeyType, PhysicalTypeValue]: """ A mapping of names to physical types for each coordinate. """ @abc.abstractmethod - def __getitem__(self, item: str): + def __getitem__(self, item: CoordKeyType) -> Any: """ Indexing the object by name should return the coordinate object. """ @abc.abstractmethod - def __iter__(self): + def __iter__(self) -> Iterator[CoordKeyType]: """ Iterate over the collection. """ @abc.abstractmethod - def __len__(self): + def __len__(self) -> int: """ Establish the length of the collection. """ @@ -83,13 +93,13 @@ def __len__(self): class GlobalCoords(GlobalCoordsABC): # Docstring in GlobalCoordsABC - def __init__(self, ndcube=None): + def __init__(self, ndcube: "NDCubeABC | None" = None) -> None: super().__init__() self._ndcube = ndcube - self._internal_coords = OrderedDict() + self._internal_coords: OrderedDict[CoordKeyType, tuple[PhysicalTypeValue, Any]] = OrderedDict() @staticmethod - def _convert_dropped_to_internal(dropped_dimensions): + def _convert_dropped_to_internal(dropped_dimensions: dict[str, Any]) -> dict[CoordKeyType, tuple[PhysicalTypeValue, Any]]: """ Convert the `~astropy.wcs.wcsapi.SlicedLowLevelWCS` style ``dropped_world_dimensions`` dictionary to the GlobalCoords internal @@ -98,7 +108,7 @@ def _convert_dropped_to_internal(dropped_dimensions): # Most of this method is adapted from # astropy.wcs.wcsapi.high_level_wcs.HighLevelWCSMixin.pixel_to_world - new_internal_coords = {} + new_internal_coords: dict[CoordKeyType, tuple[PhysicalTypeValue, Any]] = {} world = dropped_dimensions.pop("value") components = dropped_dimensions.pop("world_axis_object_components") @@ -111,8 +121,8 @@ def _convert_dropped_to_internal(dropped_dimensions): classes_new[key] = deserialize_class(value, construct=False) classes = classes_new - args = defaultdict(list) - kwargs = defaultdict(dict) + args: defaultdict[str, list[Any]] = defaultdict(list) + kwargs: defaultdict[str, dict[str, Any]] = defaultdict(dict) for i, (key, attr, _) in enumerate(components): if isinstance(attr, str): @@ -153,7 +163,7 @@ def _convert_dropped_to_internal(dropped_dimensions): return new_internal_coords @property - def _all_coords(self): + def _all_coords(self) -> dict[CoordKeyType, tuple[PhysicalTypeValue, Any]]: """ A dynamic dictionary of all global coordinates, stored here or derived from the ndcube object. @@ -161,7 +171,7 @@ def _all_coords(self): if self._ndcube is None: return self._internal_coords - all_coords = {**self._internal_coords} + all_coords: dict[CoordKeyType, tuple[PhysicalTypeValue, Any]] = {**self._internal_coords} if hasattr(self._ndcube.wcs.low_level_wcs, "dropped_world_dimensions"): dropped_world = copy.deepcopy(self._ndcube.wcs.low_level_wcs.dropped_world_dimensions) @@ -175,7 +185,7 @@ def _all_coords(self): return all_coords - def add(self, name, physical_type, coord): + def add(self, name: str, physical_type: str | None, coord: Any) -> None: # Docstring in GlobalCoordsABC if name in self._internal_coords.keys(): raise ValueError("coordinate with same name already exists: " @@ -186,16 +196,16 @@ def add(self, name, physical_type, coord): self._internal_coords[name] = (physical_type, coord) - def remove(self, name): + def remove(self, name: str) -> None: # Docstring in GlobalCoordsABC del self._internal_coords[name] @property - def physical_types(self): + def physical_types(self) -> dict[CoordKeyType, PhysicalTypeValue]: # Docstring in GlobalCoordsABC return {name: value[0] for name, value in self._all_coords.items()} - def filter_by_physical_type(self, physical_type): + def filter_by_physical_type(self, physical_type: str) -> "GlobalCoords": """ Filter this object to coordinates with a given physical type. @@ -210,10 +220,10 @@ def filter_by_physical_type(self, physical_type): A new object storing just the coordinates with the given physical type. """ gc = GlobalCoords() - gc._internal_coords = dict(filter(lambda x: x[1][0] == physical_type, self._all_coords.items())) + gc._internal_coords = OrderedDict(filter(lambda x: x[1][0] == physical_type, self._all_coords.items())) return gc - def __getitem__(self, item): + def __getitem__(self, item: CoordKeyType) -> Any: # Docstring in GlobalCoordsABC if item not in self._all_coords: for key, value in self._all_coords.items(): @@ -222,15 +232,15 @@ def __getitem__(self, item): return self._all_coords[item][1] - def __iter__(self): + def __iter__(self) -> Iterator[CoordKeyType]: # Docstring in GlobalCoordsABC return iter(self._all_coords) - def __len__(self): + def __len__(self) -> int: # Docstring in GlobalCoordsABC return len(self._all_coords) - def __str__(self): + def __str__(self) -> str: classname = self.__class__.__name__ elements = [f"{name} {[ptype]}:\n{coord!r}" for (name, coord), ptype in zip(self.items(), self.physical_types.values())] @@ -242,5 +252,5 @@ def __str__(self): return f"{classname}({joiner.join(elements)})" - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self!s}" diff --git a/ndcube/meta.py b/ndcube/meta.py index 5a10b113c..787771102 100644 --- a/ndcube/meta.py +++ b/ndcube/meta.py @@ -3,13 +3,24 @@ import numbers import collections.abc from types import MappingProxyType +from typing import Any, cast import numpy as np +import numpy.typing as npt __all__ = ["NDMeta", "NDMetaABC"] +Axes = npt.NDArray[np.integer[Any]] +# The axis/axes with which a piece of metadata is associated, either as +# supplied by a caller (a scalar int, a sequence of them, or an array of +# them) or as stored internally once sanitized (always an array of ints). +AxisSpec = int | collections.abc.Sequence[int] | Axes +# A data shape, either as stored internally (an array of ints) or as an +# arbitrary sequence of ints handed in by a caller. +ShapeLike = collections.abc.Sequence[int] | Axes -class NDMetaABC(collections.abc.Mapping): + +class NDMetaABC(collections.abc.Mapping[str, Any]): """ A sliceable object for storing metadata. @@ -35,6 +46,12 @@ class NDMetaABC(collections.abc.Mapping): number of associated axes (axis-aligned), or same shape as the associated data array's axes (grid-aligned). + data_shape: sequence of `int`, optional + The shape of the data array associated with the metadata, used to + validate grid-aligned metadata. If not given, is derived from the + shapes of any grid-aligned metadata already present, and otherwise + left empty until set explicitly. + Notes ----- **Axis-aware Metadata** @@ -76,7 +93,7 @@ class NDMetaABC(collections.abc.Mapping): @property @abc.abstractmethod - def axes(self): + def axes(self) -> collections.abc.Mapping[str, Axes]: """ Mapping from metadata keys to axes with which they are associated. @@ -85,7 +102,7 @@ def axes(self): @property @abc.abstractmethod - def key_comments(self): + def key_comments(self) -> collections.abc.Mapping[str, str]: """ Mapping from metadata keys to associated comments. @@ -94,13 +111,27 @@ def key_comments(self): @property @abc.abstractmethod - def data_shape(self): + def data_shape(self) -> Axes: """ The shape of the data with which the metadata is associated. """ + @data_shape.setter @abc.abstractmethod - def add(self, name, value, key_comment=None, axes=None, overwrite=False): + def data_shape(self, new_shape: collections.abc.Sequence[int]) -> None: + """ + Set the shape of the data with which the metadata is associated. + """ + + @abc.abstractmethod + def add( + self, + name: str, + value: Any, + key_comment: str | None = None, + axes: AxisSpec | None = None, + overwrite: bool = False, + ) -> None: """ Add a new piece of metadata to instance. @@ -126,7 +157,7 @@ def add(self, name, value, key_comment=None, axes=None, overwrite=False): @property @abc.abstractmethod - def slice(self): + def slice(self) -> "_NDMetaSlicer": """ A helper class which, when sliced, returns a new NDMeta with axis- and grid-aligned metadata sliced. @@ -136,40 +167,45 @@ def slice(self): """ @abc.abstractmethod - def rebin(self, rebinned_axes, new_shape): + def rebin(self, bin_shape: collections.abc.Sequence[int]) -> "NDMetaABC": """ Adjusts grid-aware metadata to stay consistent with rebinned data. """ -class NDMeta(dict, NDMetaABC): +class NDMeta(dict[str, Any], NDMetaABC): # Docstring in ABC __ndcube_can_slice__ = True __ndcube_can_rebin__ = True - def __init__(self, meta=None, key_comments=None, axes=None, data_shape=None): - self._original_meta = meta - if data_shape is None: - self._data_shape = np.array([], dtype=int) - else: - self._data_shape = np.asarray(data_shape).astype(int) + def __init__( + self, + meta: collections.abc.Mapping[str, Any] | None = None, + key_comments: collections.abc.Mapping[str, str] | None = None, + axes: collections.abc.Mapping[str, AxisSpec] | None = None, + data_shape: collections.abc.Sequence[int] | None = None, + ) -> None: + self._data_shape: Axes = ( + np.array([], dtype=int) if data_shape is None else np.asarray(data_shape).astype(int) + ) if meta is None: meta = {} + self._original_meta: collections.abc.Mapping[str, Any] = meta super().__init__(meta.items()) meta_keys = meta.keys() if key_comments is None: - self._key_comments = {} + self._key_comments: dict[str, str] = {} else: if not set(key_comments.keys()).issubset(set(meta_keys)): raise ValueError( "All comments must correspond to a value in meta under the same key." ) - self._key_comments = key_comments + self._key_comments = dict(key_comments) if axes is None: - self._axes = {} + self._axes: dict[str, Axes] = {} else: axes = dict(axes) if not set(axes.keys()).issubset(set(meta_keys)): @@ -178,60 +214,61 @@ def __init__(self, meta=None, key_comments=None, axes=None, data_shape=None): self._axes = {key:self._sanitize_axis_value(axis, meta[key], key) for key, axis in axes.items()} - def _sanitize_axis_value(self, axis, value, key): + def _sanitize_axis_value(self, axis: AxisSpec, value: Any, key: str) -> Axes: axis_err_msg = ("Values in axes must be an integer or iterable of integers giving " f"the data axis/axes associated with the metadata. axis = {axis}.") if isinstance(axis, numbers.Integral): - axis = (axis,) - if len(axis) == 0: + axis_seq: collections.abc.Sequence[int] | Axes = (int(axis),) + else: + axis_seq = cast("collections.abc.Sequence[int] | Axes", axis) + if len(axis_seq) == 0: raise ValueError(axis_err_msg) - # Verify each entry in axes is an iterable of ints or a scalar. - if not (isinstance(axis, collections.abc.Iterable) - and all(isinstance(i, numbers.Integral) for i in axis)): + # Verify each entry in axes is an integer. + if not all(isinstance(i, numbers.Integral) for i in axis_seq): raise ValueError(axis_err_msg) # If metadata's axis/axes include axis beyond current data shape, extend it. data_shape = self.data_shape - if max(axis) >= len(data_shape): + if max(axis_seq) >= len(data_shape): data_shape = np.concatenate((data_shape, - np.zeros(max(axis) + 1 - len(data_shape), dtype=int))) + np.zeros(max(axis_seq) + 1 - len(data_shape), dtype=int))) # Check whether metadata is compatible with data shape based on shapes # of metadata already present. - axis = np.asarray(axis) + axis_arr: Axes = np.asarray(axis_seq) if _not_scalar(value): - axis_shape = data_shape[axis] + axis_shape = data_shape[axis_arr] if not _is_axis_aligned(value, axis_shape): # If metadata corresponds to previously unconstrained axis, update data_shape. idx0 = axis_shape == 0 if idx0.any(): axis_shape[idx0] = np.array(_get_metadata_shape(value))[idx0] - data_shape[axis] = axis_shape + data_shape[axis_arr] = axis_shape # Confirm metadata is compatible with data shape. if not _is_grid_aligned(value, axis_shape): raise ValueError( - f"{key} must have same shape {tuple(data_shape[axis])} " - f"as its associated axes {axis}, ", - f"or same length as number of associated axes ({len(axis)}). " + f"{key} must have same shape {tuple(data_shape[axis_arr])} " + f"as its associated axes {axis_arr}, ", + f"or same length as number of associated axes ({len(axis_arr)}). " f"Has shape {value.shape if hasattr(value, 'shape') else len(value)}") - elif len(axis) != 1: + elif len(axis_arr) != 1: raise ValueError("Scalar and str metadata can only be assigned to one axis. " - f"key = {key}; value = {value}; axes = {axis}") + f"key = {key}; value = {value}; axes = {axis_arr}") self._data_shape = data_shape - return axis + return axis_arr @property - def key_comments(self): + def key_comments(self) -> dict[str, str]: return self._key_comments @property - def axes(self): + def axes(self) -> dict[str, Axes]: return self._axes @property - def data_shape(self): + def data_shape(self) -> Axes: return self._data_shape @data_shape.setter - def data_shape(self, new_shape): + def data_shape(self, new_shape: collections.abc.Sequence[int]) -> None: """ Set data shape to new shape. @@ -242,24 +279,31 @@ def data_shape(self, new_shape): new_shape: array-like The new shape of the data. Elements must of of type `int`. """ - new_shape = np.round(new_shape).astype(int) - if (new_shape < 0).any(): + shape_arr: Axes = np.round(new_shape).astype(int) + if (shape_arr < 0).any(): raise ValueError("new_shape cannot include negative numbers.") # Confirm input shape agrees with shapes of pre-existing metadata. old_shape = self.data_shape - if len(new_shape) != len(old_shape) and len(self._axes) > 0: + if len(shape_arr) != len(old_shape) and len(self._axes) > 0: n_meta_axes = max([ax.max() for ax in self._axes.values()]) + 1 old_shape = np.zeros(n_meta_axes, dtype=int) for key, ax in self._axes.items(): old_shape[ax] = np.asarray(self[key].shape) # Axes of length 0 are deemed to be of unknown length, and so do not have to match. idx, = np.where(old_shape > 0) - if len(idx) > 0 and (old_shape[idx] != new_shape[idx]).any(): + if len(idx) > 0 and (old_shape[idx] != shape_arr[idx]).any(): raise ValueError("new_shape not compatible with pre-existing metadata. " - f"old shape = {old_shape}, new_shape = {new_shape}") - self._data_shape = new_shape - - def add(self, name, value, key_comment=None, axes=None, overwrite=False): + f"old shape = {old_shape}, new_shape = {shape_arr}") + self._data_shape = shape_arr + + def add( + self, + name: str, + value: Any, + key_comment: str | None = None, + axes: AxisSpec | None = None, + overwrite: bool = False, + ) -> None: # Docstring in ABC. if name in self.keys() and overwrite is not True: raise KeyError(f"'{name}' already exists. " @@ -287,14 +331,14 @@ def add(self, name, value, key_comment=None, axes=None, overwrite=False): # This must be done after updating self._axes otherwise it may error. self.__setitem__(name, value) - def __delitem__(self, name): + def __delitem__(self, name: str) -> None: if name in self._key_comments: del self._key_comments[name] if name in self._axes: del self._axes[name] super().__delitem__(name) - def __setitem__(self, key, val): + def __setitem__(self, key: str, val: Any) -> None: axis = self.axes.get(key, None) if axis is not None: if _not_scalar(val): @@ -309,15 +353,15 @@ def __setitem__(self, key, val): super().__setitem__(key, val) @property - def original_meta(self): + def original_meta(self) -> MappingProxyType[str, Any]: return MappingProxyType(self._original_meta) @property - def slice(self): + def slice(self) -> "_NDMetaSlicer": # Docstring in ABC. return _NDMetaSlicer(self) - def rebin(self, bin_shape): + def rebin(self, bin_shape: collections.abc.Sequence[int]) -> "NDMeta": """ Adjusts axis-aware metadata to stay consistent with a rebinned `~ndcube.NDCube`. @@ -332,21 +376,21 @@ def rebin(self, bin_shape): The number of pixels in a bin in each dimension. """ # Sanitize input - bin_shape = np.round(bin_shape).astype(int) + bin_shape_arr: Axes = np.round(bin_shape).astype(int) data_shape = self.data_shape - bin_shape = bin_shape[:len(data_shape)] # Drop info on axes not defined by NDMeta. - if (np.mod(data_shape, bin_shape) != 0).any(): + bin_shape_arr = bin_shape_arr[:len(data_shape)] # Drop info on axes not defined by NDMeta. + if (np.mod(data_shape, bin_shape_arr) != 0).any(): raise ValueError("bin_shape must be integer factors of their associated axes.") # Remove axis-awareness from grid-aligned metadata associated with rebinned axes. - rebinned_axes = set(np.where(bin_shape != 1)[0]) + rebinned_axes = set(np.where(bin_shape_arr != 1)[0]) new_meta = copy.deepcopy(self) - null_set = set() + null_set: set[np.integer[Any]] = set() for name, axes in self.axes.items(): if (_is_grid_aligned(self[name], data_shape[axes]) and set(axes).intersection(rebinned_axes) != null_set): del new_meta._axes[name] # Update data shape. - new_meta._data_shape = new_meta._data_shape // bin_shape + new_meta._data_shape = new_meta._data_shape // bin_shape_arr return new_meta @@ -356,13 +400,13 @@ class _NDMetaSlicer: Parameters ---------- - meta: `NDMetaABC` + meta: `NDMeta` The metadata object to slice. """ - def __init__(self, meta): + def __init__(self, meta: "NDMeta") -> None: self.meta = meta - def __getitem__(self, item): + def __getitem__(self, item: Any) -> "NDMeta": data_shape = self.meta.data_shape if len(data_shape) == 0: raise TypeError("NDMeta object does not have a shape and so cannot be sliced.") @@ -407,7 +451,7 @@ def __getitem__(self, item): raise TypeError("Unrecognized slice type. " "Must be an int, slice and tuple of the same.") kept_axes = np.invert(dropped_axes) - new_meta._data_shape = new_shape[kept_axes] + new_meta._data_shape = new_shape[kept_axes] # pyright: ignore[reportPrivateUsage] # Slice all metadata associated with axes. for key, value in self.meta.items(): @@ -417,7 +461,7 @@ def __getitem__(self, item): # Calculate new axis indices. new_axis = np.asarray(list( set(axis).intersection(set(np.arange(naxes)[kept_axes])) - )) + ), dtype=int) if len(new_axis) == 0: new_axis = None else: @@ -460,7 +504,7 @@ def __getitem__(self, item): return new_meta -def _not_scalar(value): +def _not_scalar(value: Any) -> bool: return ( ( hasattr(value, "shape") @@ -472,14 +516,14 @@ def _not_scalar(value): )) -def _is_scalar(value): +def _is_scalar(value: Any) -> bool: return not _not_scalar(value) -def _get_metadata_shape(value): +def _get_metadata_shape(value: Any) -> tuple[int, ...]: return value.shape if hasattr(value, "shape") else (len(value),) -def _is_grid_aligned(value, axis_shape): +def _is_grid_aligned(value: Any, axis_shape: ShapeLike) -> bool: if _is_scalar(value): return False value_shape = _get_metadata_shape(value) @@ -488,6 +532,6 @@ def _is_grid_aligned(value, axis_shape): return True -def _is_axis_aligned(value, axis_shape): +def _is_axis_aligned(value: Any, axis_shape: ShapeLike) -> bool: len_value = len(value) if _not_scalar(value) else 1 return not _is_grid_aligned(value, axis_shape) and len_value == len(axis_shape) diff --git a/ndcube/mixins/ndslicing.py b/ndcube/mixins/ndslicing.py index 9f1d93378..4264b310d 100644 --- a/ndcube/mixins/ndslicing.py +++ b/ndcube/mixins/ndslicing.py @@ -1,15 +1,29 @@ +from typing import TYPE_CHECKING, Any + from astropy.nddata.mixins.ndslicing import NDSlicingMixin from astropy.wcs.wcsapi.wrappers.sliced_wcs import sanitize_slices +if TYPE_CHECKING: + from ndcube.extra_coords.extra_coords import ExtraCoordsABC + from ndcube.global_coords import GlobalCoordsABC + __all__ = ['NDCubeSlicingMixin'] -class NDCubeSlicingMixin(NDSlicingMixin): +class NDCubeSlicingMixin(NDSlicingMixin): # type: ignore[misc] # Inherit docstring from parent class __doc__ = NDSlicingMixin.__doc__ - def __getitem__(self, item): + # Attributes expected to be provided by the class this is mixed into + # (e.g. NDCubeBase, which combines this with astropy.nddata.NDData and NDCubeABC). + if TYPE_CHECKING: + meta: Any + shape: tuple[int, ...] + global_coords: GlobalCoordsABC + extra_coords: ExtraCoordsABC + + def __getitem__(self, item: Any) -> Any: """ Override the parent class method to explicitly catch `None` indices. @@ -23,6 +37,7 @@ def __getitem__(self, item): # This is to prevent the shapes of the data and metadata getting out of # sync part way through the slicing process. meta_is_sliceable = False + meta: Any = None if hasattr(self.meta, "__ndcube_can_slice__") and self.meta.__ndcube_can_slice__: meta_is_sliceable = True meta = self.meta @@ -46,7 +61,7 @@ def __getitem__(self, item): self.meta = meta # Add unsliced meta back onto unsliced cube. # Add sliced coords back onto sliced cube. - sliced_cube._global_coords._internal_coords = self.global_coords._internal_coords + sliced_cube._global_coords._internal_coords = self.global_coords._internal_coords # type: ignore[attr-defined] # pyright: ignore[reportAttributeAccessIssue] sliced_cube._extra_coords = self.extra_coords[item] # If metadata sliceable, slice and add back onto sliced cube. diff --git a/ndcube/ndcollection.py b/ndcube/ndcollection.py index 68a06babd..82ac7eff7 100644 --- a/ndcube/ndcollection.py +++ b/ndcube/ndcollection.py @@ -2,6 +2,7 @@ import numbers import textwrap import collections.abc +from typing import Any import numpy as np @@ -11,7 +12,7 @@ __all__ = ["NDCollection"] -class NDCollection(dict): +class NDCollection(dict[Any, Any]): """ Class for holding and manipulating an unordered collection of NDCubes or NDCubeSequences. @@ -49,7 +50,7 @@ class NDCollection(dict): axis 1 of cube0 is aligned with axis 1 of cube1. """ - def __init__(self, key_data_pairs, aligned_axes=None, meta=None, **kwargs): + def __init__(self, key_data_pairs: Any, aligned_axes: Any = None, meta: Any = None, **kwargs: Any) -> None: if isinstance(key_data_pairs, collections.abc.Mapping): key_data_pairs = tuple(key_data_pairs.items()) for key, _ in key_data_pairs: @@ -64,11 +65,12 @@ def __init__(self, key_data_pairs, aligned_axes=None, meta=None, **kwargs): # Convert aligned axes to required format. sanitize_inputs = kwargs.pop("sanitize_inputs", True) + keys: tuple[Any, ...] = () if aligned_axes is not None: keys, data = zip(*key_data_pairs) # Sanitize aligned axes unless hidden kwarg indicates not to. if sanitize_inputs: - aligned_axes = collection_utils._sanitize_aligned_axes(keys, data, aligned_axes) + aligned_axes = collection_utils._sanitize_aligned_axes(keys, data, aligned_axes) # pyright: ignore[reportPrivateUsage] else: aligned_axes = dict(zip(keys, aligned_axes)) if kwargs: @@ -80,20 +82,20 @@ def __init__(self, key_data_pairs, aligned_axes=None, meta=None, **kwargs): if self._aligned_axes is None: self.n_aligned_axes = 0 else: - self.n_aligned_axes = len(self.aligned_axes[keys[0]]) + self.n_aligned_axes = len(self._aligned_axes[keys[0]]) @property - def aligned_axes(self): + def aligned_axes(self) -> dict[Any, tuple[Any, ...]] | None: """ The axes of each array that are aligned in numpy order. """ return self._aligned_axes @property - def _first_key(self): + def _first_key(self) -> Any: return list(self.keys())[0] - def __str__(self): + def __str__(self) -> str: return (textwrap.dedent(f"""\ NDCollection ------------ @@ -102,11 +104,11 @@ def __str__(self): Aligned dimensions: {self.aligned_dimensions} Aligned physical types: {self.aligned_axis_physical_types}""")) - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self!s}" @property - def aligned_dimensions(self): + def aligned_dimensions(self) -> Any: """ The lengths of all aligned axes. @@ -119,7 +121,7 @@ def aligned_dimensions(self): return None @property - def aligned_axis_physical_types(self): + def aligned_axis_physical_types(self) -> list[tuple[Any, ...]] | None: """ The physical types common to all members that are associated with each aligned axis. @@ -138,7 +140,7 @@ def aligned_axis_physical_types(self): return [tuple(set.intersection(*[set(cube_types[i]) for cube_types in collection_types])) for i in range(self.n_aligned_axes)] - def __getitem__(self, item): + def __getitem__(self, item: Any) -> Any: # There are two ways to slice: # by key or sequence of keys, i.e. slice out given cubes in the collection, or # by typical python numeric slicing API, @@ -162,7 +164,8 @@ def __getitem__(self, item): if item_is_strings: new_data = [self[_item] for _item in item] new_keys = item - new_aligned_axes = tuple([self.aligned_axes[item_] for item_ in item]) + new_aligned_axes = (tuple([self.aligned_axes[item_] for item_ in item]) + if self.aligned_axes is not None else None) new_meta = copy.deepcopy(self.meta) # Else, the item is assumed to be a typical slicing item. @@ -189,9 +192,11 @@ def __getitem__(self, item): if item < 0: sanitized_item = int(self.aligned_dimensions[0] + item) elif isinstance(item, slice): - if (item.start is not None and item.start < 0) or (item.stop is not None and item.stop < 0): - new_start = aligned_shape[0] + item.start if item.start < 0 else item.start - new_stop = aligned_shape[0] + item.stop if item.stop < 0 else item.stop + start_neg = item.start is not None and item.start < 0 + stop_neg = item.stop is not None and item.stop < 0 + if start_neg or stop_neg: + new_start = aligned_shape[0] + item.start if start_neg else item.start + new_stop = aligned_shape[0] + item.stop if stop_neg else item.stop sanitized_item = slice(new_start, new_stop) else: sanitized_item = list(sanitized_item) @@ -199,9 +204,11 @@ def __getitem__(self, item): if isinstance(ax_it, numbers.Integral) and ax_it < 0: sanitized_item[i] = aligned_shape[i] + ax_it elif isinstance(ax_it, slice): - if (ax_it.start is not None and ax_it.start < 0) or (ax_it.stop is not None and ax_it.stop < 0): - new_start = aligned_shape[i] + ax_it.start if ax_it.start < 0 else ax_it.start - new_stop = aligned_shape[i] + ax_it.stop if ax_it.stop < 0 else ax_it.stop + start_neg = ax_it.start is not None and ax_it.start < 0 + stop_neg = ax_it.stop is not None and ax_it.stop < 0 + if start_neg or stop_neg: + new_start = aligned_shape[i] + ax_it.start if start_neg else ax_it.start + new_stop = aligned_shape[i] + ax_it.stop if stop_neg else ax_it.stop sanitized_item[i] = slice(new_start, new_stop) sanitized_item = tuple(sanitized_item) # Use sanitized item to slice meta. @@ -216,15 +223,19 @@ def __getitem__(self, item): sanitize_inputs=False ) - def _generate_collection_getitems(self, item): + def _generate_collection_getitems(self, item: Any) -> tuple[list[list[Any]], Any]: # There are 3 supported cases of the slice item: int, slice, tuple of ints and/or slices. # Compile appropriate slice items for each cube in the collection and # and drop any aligned axes that are sliced out. + # Only called from __getitem__ after it has confirmed aligned_axes is not None. + if self.aligned_axes is None: + raise RuntimeError("aligned_axes must not be None to generate collection getitems.") + # First, define empty lists of slice items to be applied to each cube in collection. - collection_items = [[slice(None)] * len(self[key].shape) for key in self] + collection_items: list[list[Any]] = [[slice(None)] * len(self[key].shape) for key in self] # Define empty list to hold aligned axes dropped by the slicing. - drop_aligned_axes_indices = [] + drop_aligned_axes_indices: list[Any] = [] # Case 1: int # First aligned axis is dropped. @@ -255,30 +266,30 @@ def _generate_collection_getitems(self, item): collection_items[j][self.aligned_axes[key][i]] = axis_item else: - raise TypeError(f"Unsupported slicing type: {axis_item}") + raise TypeError(f"Unsupported slicing type: {item}") # Use indices of dropped axes determine above to update aligned_axes # by removing any that have been dropped. - drop_aligned_axes_indices = np.array(drop_aligned_axes_indices) - new_aligned_axes = collection_utils._update_aligned_axes( - drop_aligned_axes_indices, self.aligned_axes, self._first_key) + drop_aligned_axes_indices_arr = np.array(drop_aligned_axes_indices) + new_aligned_axes = collection_utils._update_aligned_axes( # pyright: ignore[reportPrivateUsage] + drop_aligned_axes_indices_arr, self.aligned_axes, self._first_key) return collection_items, new_aligned_axes - def copy(self): + def copy(self) -> "NDCollection": # Aligned axes is not a required parameter and may be None aligned_axes = None if self.aligned_axes is None else tuple(self.aligned_axes.values()) return self.__class__(self.items(), aligned_axes, meta=self.meta, sanitize_inputs=False) - def setdefault(self): + def setdefault(self) -> Any: # type: ignore[override] """Not supported by `~ndcube.NDCollection`""" raise NotImplementedError("NDCollection does not support setdefault.") - def popitem(self): + def popitem(self) -> Any: """Not supported by `~ndcube.NDCollection`""" raise NotImplementedError("NDCollection does not support popitem.") - def pop(self, key): + def pop(self, key: Any) -> Any: # type: ignore[override] """ Remove a member from the `~ndcube.NDCollection` and return it. @@ -295,7 +306,7 @@ def pop(self, key): self.aligned_axes.pop(key) return popped_cube - def update(self, *args): + def update(self, *args: Any) -> None: # type: ignore[override] """ Merges a new collection with current one replacing objects with common keys. @@ -306,7 +317,7 @@ def update(self, *args): if len(args) == 2: key_data_pairs = args[0] new_keys, new_data = zip(*key_data_pairs) - new_aligned_axes = collection_utils._sanitize_aligned_axes(new_keys, new_data, args[1]) + new_aligned_axes = collection_utils._sanitize_aligned_axes(new_keys, new_data, args[1]) # pyright: ignore[reportPrivateUsage] else: # If one arg given, input must be NDCollection. collection = args[0] new_keys = list(collection.keys()) @@ -325,12 +336,15 @@ def update(self, *args): # Update collection super().update(key_data_pairs) if first_old_aligned_axes is not None: # since the above assertion passed, if one aligned axes is not None, both are not None + if self.aligned_axes is None or new_aligned_axes is None: + raise RuntimeError("aligned_axes must not be None here; assert_aligned_axes_compatible should have raised.") self.aligned_axes.update(new_aligned_axes) - def __delitem__(self, key): + def __delitem__(self, key: Any) -> None: super().__delitem__(key) - self.aligned_axes.__delitem__(key) + if self.aligned_axes is not None: + self.aligned_axes.__delitem__(key) - def __setitem__(self, key, value): + def __setitem__(self, key: Any, value: Any) -> None: raise NotImplementedError("NDCollection does not support __setitem__. " "Use NDCollection.update instead") diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index ee1b0909a..e77bc7e07 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -18,7 +18,7 @@ try: # Import sunpy coordinates if available to register the frames and WCS functions with astropy - import sunpy.coordinates # NOQA + import sunpy.coordinates # NOQA # pyright: ignore[reportUnusedImport] except ImportError: pass @@ -39,17 +39,17 @@ __all__ = ['NDCubeABC', 'NDCubeLinkedDescriptor'] # Create mapping to masked array types based on data array type for use in analysis methods. -ARRAY_MASK_MAP = {} +ARRAY_MASK_MAP: dict[type, Any] = {} ARRAY_MASK_MAP[np.ndarray] = np.ma.masked_array _NUMPY_COPY_IF_NEEDED = False if np.__version__.startswith("1.") else None try: import dask.array - ARRAY_MASK_MAP[dask.array.core.Array] = dask.array.ma.masked_array + ARRAY_MASK_MAP[dask.array.core.Array] = dask.array.ma.masked_array # pyright: ignore[reportAttributeAccessIssue, reportPrivateImportUsage, reportFunctionMemberAccess] except ImportError: pass -class NDCubeABC(astropy.nddata.NDDataBase): +class NDCubeABC(astropy.nddata.NDDataBase): # type: ignore[misc] @property @abc.abstractmethod @@ -114,7 +114,7 @@ def axis_world_coords(self, If `True` then instead of returning the coordinates at the centers of the pixels, the coordinates at the pixel corners will be returned. This increases the size of the output by 1 in all dimensions as all corners are returned. - wcs: `astropy.wcs.wcsapi.BaseHighLevelWCS`, optional + wcs: `astropy.wcs.wcsapi.BaseHighLevelWCS` or `~ndcube.ExtraCoordsABC`, optional The WCS object to used to calculate the world coordinates. Although technically this can be any valid WCS, it will typically be ``self.wcs``, ``self.extra_coords``, or ``self.combined_wcs`` combining both @@ -293,11 +293,11 @@ class NDCubeLinkedDescriptor: A descriptor which gives the property a reference to the cube to which it is attached. """ - def __init__(self, default_type): + def __init__(self, default_type: type) -> None: self._default_type = default_type - self._property_name = None + self._property_name: str | None = None - def __set_name__(self, owner, name): + def __set_name__(self, owner: type, name: str) -> None: """ This function is called when the class the descriptor is attached to is initialized. @@ -310,18 +310,18 @@ def __set_name__(self, owner, name): # the data is stored. self._attribute_name = f"_{name}" - def __get__(self, obj, objtype=None): + def __get__(self, obj: Any, objtype: type | None = None) -> Any: if obj is None: return None - if getattr(obj, self._attribute_name, None) is None and self._default_type is not None: + if getattr(obj, self._attribute_name, None) is None and self._default_type is not None: # pyright: ignore[reportUnnecessaryComparison] self.__set__(obj, self._default_type) return getattr(obj, self._attribute_name) - def __set__(self, obj, value): + def __set__(self, obj: Any, value: Any) -> None: if isinstance(value, self._default_type): - value._ndcube = obj + value._ndcube = obj # type: ignore[attr-defined] elif issubclass(value, self._default_type): value = value(obj) else: @@ -332,7 +332,7 @@ def __set__(self, obj, value): setattr(obj, self._attribute_name, value) -class NDCubeBase(NDCubeABC, astropy.nddata.NDData, NDCubeSlicingMixin): +class NDCubeBase(NDCubeABC, astropy.nddata.NDData, NDCubeSlicingMixin): # type: ignore[misc] """ Class representing N-D data described by a single array and set of WCS transformations. @@ -377,9 +377,13 @@ class NDCubeBase(NDCubeABC, astropy.nddata.NDData, NDCubeSlicingMixin): # Instances of Extra and Global coords are managed through descriptors _extra_coords = NDCubeLinkedDescriptor(ExtraCoords) _global_coords = NDCubeLinkedDescriptor(GlobalCoords) + # Inherited from astropy.nddata.NDData; declared here so mypy can resolve its type + # (NDData is unstubbed, and self.mask is both read and reassigned in this class). + mask: Any - def __init__(self, data, wcs=None, uncertainty=None, mask=None, meta=None, - unit=None, copy=False, psf=None, *, extra_coords=None, global_coords=None, **kwargs): + def __init__(self, data: Any, wcs: Any = None, uncertainty: Any = None, mask: Any = None, meta: Any = None, + unit: Any = None, copy: bool = False, psf: Any = None, *, extra_coords: Any = None, + global_coords: Any = None, **kwargs: Any) -> None: super().__init__(data, wcs=wcs, uncertainty=uncertainty, mask=mask, meta=meta, unit=unit, copy=copy, psf=psf, **kwargs) @@ -409,7 +413,7 @@ def __init__(self, data, wcs=None, uncertainty=None, mask=None, meta=None, self.meta.data_shape = self.shape @property - def data(self): + def data(self) -> Any: """ `~numpy.ndarray` - like The stored dataset. @@ -425,7 +429,7 @@ def data(self): return super().data @data.setter - def data(self, value): + def data(self, value: Any) -> None: # In an array-agnostic way check the shape is the same if not hasattr(value, "shape") or value.shape != self.data.shape: raise TypeError(f"Can only set data with an array-like object of the same shape ({self.data.shape})") @@ -446,37 +450,37 @@ def data(self, value): self._data = value @property - def extra_coords(self): + def extra_coords(self) -> ExtraCoordsABC: # type: ignore[override] # Docstring in NDCubeABC. - return self._extra_coords + return self._extra_coords # type: ignore[no-any-return] @property - def global_coords(self): + def global_coords(self) -> GlobalCoordsABC: # type: ignore[override] # Docstring in NDCubeABC. - return self._global_coords + return self._global_coords # type: ignore[no-any-return] @property - def combined_wcs(self): + def combined_wcs(self) -> BaseHighLevelWCS: # Docstring in NDCubeABC. if not self.extra_coords.wcs: return self.wcs - mapping = list(range(self.wcs.pixel_n_dim)) + list(self.extra_coords.mapping) + mapping = tuple(range(self.wcs.pixel_n_dim)) + tuple(self.extra_coords.mapping) return HighLevelWCSWrapper( CompoundLowLevelWCS(self.wcs.low_level_wcs, self._extra_coords.wcs, mapping=mapping) ) @property - def shape(self): - return self.data.shape + def shape(self) -> tuple[int, ...]: # type: ignore[override] + return self.data.shape # type: ignore[no-any-return] @property - def dimensions(self): + def dimensions(self) -> u.Quantity: warn_deprecated("Replaced by ndcube.NDCube.shape") return u.Quantity(self.data.shape, unit=u.pix) @property - def array_axis_physical_types(self): + def array_axis_physical_types(self) -> list[tuple[str, ...]]: # Docstring in NDCubeABC. wcs = self.combined_wcs world_axis_physical_types = np.array(wcs.world_axis_physical_types) @@ -485,11 +489,12 @@ def array_axis_physical_types(self): for i in range(axis_correlation_matrix.shape[1])][::-1] @property - def quantity(self): + def quantity(self) -> u.Quantity: """Unitful representation of the NDCube data.""" return u.Quantity(self.data, self.unit, copy=_NUMPY_COPY_IF_NEEDED) - def _generate_world_coords(self, pixel_corners, wcs, *, needed_axes, units=None): + def _generate_world_coords(self, pixel_corners: bool, wcs: Any, *, needed_axes: Any, + units: bool | None = None) -> Any: """ Private method to generate world coordinates. @@ -564,7 +569,7 @@ def _generate_world_coords(self, pixel_corners, wcs, *, needed_axes, units=None) return world_coords @utils.cube.sanitize_wcs - def axis_world_coords(self, *axes, pixel_corners=False, wcs=None): + def axis_world_coords(self, *axes: int | str, pixel_corners: bool = False, wcs: Any = None) -> Any: # Docstring in NDCubeABC. if isinstance(wcs, BaseHighLevelWCS): wcs = wcs.low_level_wcs @@ -592,7 +597,7 @@ def axis_world_coords(self, *axes, pixel_corners=False, wcs=None): return tuple(axes_coords[i] for i in object_indices) @utils.cube.sanitize_wcs - def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None): + def axis_world_coords_values(self, *axes: int | str, pixel_corners: bool = False, wcs: Any = None) -> Any: # Docstring in NDCubeABC. if isinstance(wcs, BaseHighLevelWCS): wcs = wcs.low_level_wcs @@ -617,19 +622,19 @@ def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None): identifier = identifier.replace(".", "_") identifier = identifier.replace("-", "__") identifiers.append(identifier) - CoordValues = namedtuple("CoordValues", identifiers) + CoordValues = namedtuple("CoordValues", identifiers) # type: ignore[misc] return CoordValues(*axes_coords[::-1]) - def crop(self, *points, wcs=None, keepdims=False): + def crop(self, *points: Any, wcs: Any = None, keepdims: bool = False) -> "NDCubeBase": # The docstring is defined in NDCubeABC # Calculate the array slice item corresponding to bounding box and return sliced cube. item = self._get_crop_item(*points, wcs=wcs, keepdims=keepdims) - return self[item] + return self[item] # type: ignore[no-any-return] @utils.cube.sanitize_wcs - def _get_crop_item(self, *points, wcs=None, keepdims=False): + def _get_crop_item(self, *points: Any, wcs: Any = None, keepdims: bool = False) -> Any: # Sanitize inputs. - no_op, points, wcs = utils.cube.sanitize_crop_inputs(points, wcs) + no_op, sanitized_points, wcs = utils.cube.sanitize_crop_inputs(points, wcs) # Quit out early if we are no-op if no_op: return tuple([slice(None)] * wcs.pixel_n_dim) @@ -640,7 +645,7 @@ def _get_crop_item(self, *points, wcs=None, keepdims=False): if comp.count(c) > 1: comp.pop(k) classes = [wcs.world_axis_object_classes[c][0] for c in comp] - for i, point in enumerate(points): + for i, point in enumerate(sanitized_points): if len(point) != len(comp): raise ValueError(f"{len(point)} components in point {i} do not match " f"WCS with {len(comp)} components.") @@ -649,30 +654,30 @@ def _get_crop_item(self, *points, wcs=None, keepdims=False): raise TypeError(f"{type(value)} of component {j} in point {i} is " f"incompatible with WCS component {comp[j]} " f"{classes[j]}.") - return utils.cube.get_crop_item_from_points(points, wcs, False, keepdims=keepdims, + return utils.cube.get_crop_item_from_points(sanitized_points, wcs, False, keepdims=keepdims, original_shape=self.data.shape) - def crop_by_values(self, *points, units=None, wcs=None, keepdims=False): + def crop_by_values(self, *points: Any, units: Any = None, wcs: Any = None, keepdims: bool = False) -> "NDCubeBase": # The docstring is defined in NDCubeABC # Calculate the array slice item corresponding to bounding box and return sliced cube. item = self._get_crop_by_values_item(*points, units=units, wcs=wcs, keepdims=keepdims) - return self[item] + return self[item] # type: ignore[no-any-return] @utils.cube.sanitize_wcs - def _get_crop_by_values_item(self, *points, units=None, wcs=None, keepdims=False): + def _get_crop_by_values_item(self, *points: Any, units: Any = None, wcs: Any = None, keepdims: bool = False) -> Any: # Sanitize inputs. - no_op, points, wcs = utils.cube.sanitize_crop_inputs(points, wcs) + no_op, sanitized_points, wcs = utils.cube.sanitize_crop_inputs(points, wcs) # Quit out early if we are no-op if no_op: return tuple([slice(None)] * wcs.pixel_n_dim) # Convert float inputs to quantities using units. - n_coords = len(points[0]) + n_coords = len(sanitized_points[0]) if units is None: units = [None] * n_coords elif len(units) != n_coords: raise ValueError(f"Units must be None or have same length {n_coords} as corner inputs.") types_with_units = (u.Quantity, type(None)) - for i, point in enumerate(points): + for i, point in enumerate(sanitized_points): if len(point) != wcs.world_n_dim: raise ValueError(f"{len(point)} dimensions in point {i} do not match " f"WCS with {wcs.world_n_dim} world dimensions.") @@ -684,18 +689,18 @@ def _get_crop_by_values_item(self, *points, units=None, wcs=None, keepdims=False "If an element of a point is not a Quantity or None, " "the corresponding unit must be a valid astropy Unit or unit string." f"index: {i}; coord type: {type(value)}; unit: {unit}") - points[i][j] = u.Quantity(value, unit=unit) + sanitized_points[i][j] = u.Quantity(value, unit=unit) if value is not None: try: - points[i][j] = points[i][j].to(wcs.world_axis_units[j]) + sanitized_points[i][j] = sanitized_points[i][j].to(wcs.world_axis_units[j]) except UnitsError as err: - raise UnitsError(f"Unit '{points[i][j].unit}' of coordinate object {j} in point {i} is " + raise UnitsError(f"Unit '{sanitized_points[i][j].unit}' of coordinate object {j} in point {i} is " f"incompatible with WCS unit '{wcs.world_axis_units[j]}'") from err - return utils.cube.get_crop_item_from_points(points, wcs, True, keepdims=keepdims, + return utils.cube.get_crop_item_from_points(sanitized_points, wcs, True, keepdims=keepdims, original_shape=self.data.shape) - def __str__(self): + def __str__(self) -> str: return textwrap.dedent(f"""\ NDCube ------ @@ -704,10 +709,10 @@ def __str__(self): Unit: {self.unit} Data Type: {self.data.dtype}""") - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self!s}" - def explode_along_axis(self, axis): + def explode_along_axis(self, axis: int) -> "NDCubeSequence": """ Separates slices of NDCubes along a given axis into an NDCubeSequence of (N-1)DCubes. @@ -726,7 +731,7 @@ def explode_along_axis(self, axis): # To store the resultant cube result_cubes = [] # All slices are initially initialised as slice(None, None, None) - cube_slices = [slice(None, None, None)] * self.data.ndim + cube_slices: list[Any] = [slice(None, None, None)] * self.data.ndim # Slicing the cube inside result_cube for i in range(self.data.shape[axis]): # Setting the slice value to the index so that the slices are done correctly. @@ -741,11 +746,11 @@ def explode_along_axis(self, axis): return NDCubeSequence(result_cubes, meta=self.meta) def reproject_to(self, - target_wcs, - algorithm='interpolation', - shape_out=None, - return_footprint=False, - **reproject_args): + target_wcs: Any, + algorithm: str = 'interpolation', + shape_out: tuple[int, ...] | None = None, + return_footprint: bool = False, + **reproject_args: Any) -> Any: """ Reprojects the instance to the coordinates described by another WCS object. @@ -849,6 +854,7 @@ def reproject_to(self, return_footprint=return_footprint, **reproject_args) + footprint = None if return_footprint: data, footprint = data @@ -926,7 +932,7 @@ class NDCube(NDCubeBase): The type of this attribute can be changed to provide custom visualization functionality. """ - def _as_mpl_axes(self): + def _as_mpl_axes(self) -> Any: if hasattr(self.plotter, "_as_mpl_axes"): return self.plotter._as_mpl_axes() warn_user(f"The current plotter {self.plotter} does not have a '_as_mpl_axes' method. " @@ -935,9 +941,9 @@ def _as_mpl_axes(self): from ndcube.visualization.mpl_plotter import MatplotlibPlotter # noqa: PLC0415 plotter = MatplotlibPlotter(self) - return plotter._as_mpl_axes() + return plotter._as_mpl_axes() # pyright: ignore[reportPrivateUsage] - def plot(self, *args, **kwargs): + def plot(self, *args: Any, **kwargs: Any) -> Any: """ A convenience function for the plotters default ``plot()`` method. @@ -953,7 +959,7 @@ def plot(self, *args, **kwargs): return self.plotter.plot(*args, **kwargs) - def _new_instance(self, **kwargs): + def _new_instance(self, **kwargs: Any) -> "NDCube": keys = ('unit', 'wcs', 'mask', 'meta', 'uncertainty', 'psf') new_kwargs = {k: deepcopy(getattr(self, k, None)) for k in keys} # To support old versions of astropy, we need to make sure @@ -963,17 +969,17 @@ def _new_instance(self, **kwargs): # We Explicitly DO NOT deepcopy any data full_kwargs['data'] = self.data full_kwargs.update(kwargs) - new_cube = type(self)(**full_kwargs) - if self.extra_coords is not None: + new_cube = type(self)(**full_kwargs) # type: ignore[arg-type] + if self.extra_coords is not None: # pyright: ignore[reportUnnecessaryComparison] new_cube._extra_coords = deepcopy(self.extra_coords) - if self.global_coords is not None: + if self.global_coords is not None: # pyright: ignore[reportUnnecessaryComparison] new_cube._global_coords = deepcopy(self.global_coords) return new_cube - def __neg__(self): + def __neg__(self) -> "NDCube": return self._new_instance(data=-self.data) - def _arithmetic_handle_mask(self, self_mask, value_mask): + def _arithmetic_handle_mask(self, self_mask: Any, value_mask: Any) -> Any: if self_mask is None and value_mask is None: return None if self_mask is None: @@ -982,11 +988,11 @@ def _arithmetic_handle_mask(self, self_mask, value_mask): return self_mask return np.logical_or(self_mask, value_mask) - def _arithmetic_operate_with_nddata(self, operation, value): + def _arithmetic_operate_with_nddata(self, operation: str, value: Any) -> dict[str, Any]: handle_mask = self._arithmetic_handle_mask if value.wcs is not None: raise TypeError("Cannot add coordinate-aware objects to NDCubes.") - kwargs = {} + kwargs: dict[str, Any] = {} if operation == "add": # Handle units if self.unit is not None and value.unit is not None: @@ -1029,8 +1035,8 @@ def _arithmetic_operate_with_nddata(self, operation, value): return kwargs - def __add__(self, value): - kwargs = {} + def __add__(self, value: Any) -> "NDCube": + kwargs: dict[str, Any] = {} if isinstance(value, NDData): kwargs = self._arithmetic_operate_with_nddata("add", value) elif hasattr(value, 'unit'): @@ -1053,7 +1059,7 @@ def __add__(self, value): # return the new NDCube instance return self._new_instance(**kwargs) - def _combine_uncertainty(self, operation, value, result_data): + def _combine_uncertainty(self, operation: Any, value: Any, result_data: Any) -> Any: # combine the uncertainty; if self.uncertainty is not None and value.uncertainty is not None: if self.unit is not None: @@ -1068,10 +1074,10 @@ def _combine_uncertainty(self, operation, value, result_data): return value.uncertainty return None - def __radd__(self, value): + def __radd__(self, value: Any) -> "NDCube": return self.__add__(value) - def __sub__(self, value): + def __sub__(self, value: Any) -> "NDCube": if isinstance(value, NDData): new_value = copy.copy(value) new_value._data = -value.data @@ -1079,11 +1085,11 @@ def __sub__(self, value): new_value = -value return self.__add__(new_value) - def __rsub__(self, value): + def __rsub__(self, value: Any) -> "NDCube": return self.__neg__().__add__(value) - def __mul__(self, value): - kwargs = {} + def __mul__(self, value: Any) -> "NDCube": + kwargs: dict[str, Any] = {} if isinstance(value, NDData): kwargs = self._arithmetic_operate_with_nddata("multiply", value) else: @@ -1105,19 +1111,19 @@ def __mul__(self, value): if self.uncertainty is not None else None) return self._new_instance(**kwargs) - def __rmul__(self, value): + def __rmul__(self, value: Any) -> "NDCube": return self.__mul__(value) - def __truediv__(self, value): + def __truediv__(self, value: Any) -> "NDCube": if isinstance(value, NDData): kwargs = self._arithmetic_operate_with_nddata("true_divide", value) return self._new_instance(**kwargs) return self.__mul__(1/value) - def __rtruediv__(self, value): + def __rtruediv__(self, value: Any) -> "NDCube": return self.__pow__(-1).__mul__(value) - def __pow__(self, value): + def __pow__(self, value: Any) -> "NDCube": new_data = self.data ** value new_unit = self.unit if self.unit is None else self.unit ** value new_uncertainty = self.uncertainty @@ -1138,7 +1144,7 @@ def __pow__(self, value): return self._new_instance(data=new_data, unit=new_unit, uncertainty=new_uncertainty) - def to(self, new_unit, **kwargs): + def to(self, new_unit: Any, **kwargs: Any) -> "NDCube": """Convert instance to another unit. Converts the data, uncertainty and unit and returns a new instance @@ -1157,10 +1163,11 @@ def to(self, new_unit, **kwargs): A new instance with the new unit and data and uncertainties scales accordingly. """ new_unit = u.Unit(new_unit) - return self * (self.unit.to(new_unit, **kwargs) * new_unit / self.unit) + return self * (self.unit.to(new_unit, **kwargs) * new_unit / self.unit) # type: ignore[no-any-return] - def rebin(self, bin_shape, operation=np.mean, operation_ignores_mask=False, handle_mask=np.all, - propagate_uncertainties=False, new_unit=None, **kwargs): + def rebin(self, bin_shape: Any, operation: Any = np.mean, operation_ignores_mask: bool = False, + handle_mask: Any = np.all, propagate_uncertainties: bool | Any = False, + new_unit: Any = None, **kwargs: Any) -> "NDCube": """ Downsample array by combining contiguous pixels into bins. @@ -1360,7 +1367,7 @@ def my_propagate(uncertainty, data, mask, **kwargs): flat_uncertainty = flat_uncertainty.reshape(flat_shape) flat_uncertainty = type(self.uncertainty)(flat_uncertainty) if sanitized_mask is not None: - reshaped_mask = self.mask.reshape(tuple(reshape)) + reshaped_mask = self.mask.reshape(tuple(reshape)) # pyright: ignore[reportOptionalMemberAccess, reportAttributeAccessIssue] flat_mask = np.moveaxis(reshaped_mask, dummy_axes, tuple(range(naxes))) flat_mask = flat_mask.reshape(flat_shape) else: @@ -1396,7 +1403,7 @@ def my_propagate(uncertainty, data, mask, **kwargs): return new_cube - def squeeze(self, axis=None): + def squeeze(self, axis: Any = None) -> "NDCube": """ Removes all axes with a length of 1. @@ -1428,10 +1435,11 @@ def squeeze(self, axis=None): if (item == 0).all(): raise ValueError("All axes are of length 1, therefore we will not squeeze NDCube to become a scalar. " "Use `axis=` keyword to specify a subset of axes to squeeze.") - return self[tuple(item)] + return self[tuple(item)] # type: ignore[no-any-return] - def fill_masked(self, fill_value, uncertainty_fill_value=None, unmask=False, fill_in_place=False): + def fill_masked(self, fill_value: Any, uncertainty_fill_value: Any = None, unmask: bool = False, + fill_in_place: bool = False) -> "NDCube | None": """ Replaces masked data values with input value. @@ -1484,10 +1492,11 @@ def fill_masked(self, fill_value, uncertainty_fill_value=None, unmask=False, fil if not fill_in_place: # Create kwargs dictionary and return a new instance. - kwargs = {} + # new_mask is set above whenever fill_in_place is False. + kwargs: dict[str, Any] = {} kwargs['data'] = new_data kwargs['uncertainty'] = new_uncertainty - kwargs['mask'] = new_mask + kwargs['mask'] = new_mask # pyright: ignore[reportPossiblyUnboundVariable] return self._new_instance(**kwargs) if unmask: self.mask = False @@ -1495,16 +1504,16 @@ def fill_masked(self, fill_value, uncertainty_fill_value=None, unmask=False, fil def to_nddata(self, *, - data="copy", - wcs="copy", - uncertainty="copy", - mask="copy", - unit="copy", - meta="copy", - psf="copy", - nddata_type=NDData, - **kwargs, - ): + data: Any = "copy", + wcs: Any = "copy", + uncertainty: Any = "copy", + mask: Any = "copy", + unit: Any = "copy", + meta: Any = "copy", + psf: Any = "copy", + nddata_type: type = NDData, + **kwargs: Any, + ) -> Any: """ Constructs a new `~astropy.nddata.NDData` instance from this object. @@ -1537,7 +1546,7 @@ def to_nddata(self, Metadata object of new instance. Default is to use data of this instance. psf: Any, optional PSF object of new instance. Default is to use data of this instance. - nddata_type: Any, optional + nddata_type: `type`, optional The type of the returned object. Must be a subclass of `~astropy.nddata.NDData` or a class that behaves like one. Default=`~astropy.nddata.NDData`. kwargs: @@ -1657,7 +1666,7 @@ def to_nddata(self, return nddata_type(**user_kwargs) -def _create_masked_array_for_rebinning(data, mask, operation_ignores_mask): +def _create_masked_array_for_rebinning(data: Any, mask: Any, operation_ignores_mask: bool) -> tuple[Any, Any]: m = None if (mask is None or mask is False or operation_ignores_mask) else mask if m is None: return data, m diff --git a/ndcube/ndcube_sequence.py b/ndcube/ndcube_sequence.py index c32de7fb3..55a770432 100644 --- a/ndcube/ndcube_sequence.py +++ b/ndcube/ndcube_sequence.py @@ -1,6 +1,8 @@ import copy import numbers import textwrap +from typing import Any +from collections.abc import Iterator, Sequence import numpy as np @@ -34,16 +36,18 @@ class NDCubeSequenceBase: were a single cube concatenated along the common axis. """ - def __init__(self, data_list, meta=None, common_axis=None, **kwargs): + def __init__(self, data_list: Sequence[Any], meta: Any = None, common_axis: int | None = None, + **kwargs: Any) -> None: self.data = data_list self.meta = meta + self._common_axis: int | None if common_axis is not None: self._common_axis = int(common_axis) else: self._common_axis = common_axis @property - def dimensions(self): + def dimensions(self) -> tuple[u.Quantity, ...]: """ The length of each axis including the sequence axis. """ @@ -51,11 +55,11 @@ def dimensions(self): return tuple([d * u.pix for d in self._shape]) @property - def shape(self): + def shape(self) -> tuple[Any, ...]: return self._shape @property - def _shape(self): + def _shape(self) -> tuple[Any, ...]: dimensions = [len(self.data), *list(self.data[0].data.shape)] if len(dimensions) > 1: # If there is a common axis, length of cube's along it may not @@ -70,14 +74,14 @@ def _shape(self): return tuple(dimensions) @property - def array_axis_physical_types(self): + def array_axis_physical_types(self) -> list[tuple[str, ...]]: """ The physical types associated with each array axis, including the sequence axis. """ return [("meta.obs.sequence",), *self.data[0].array_axis_physical_types] @property - def cube_like_dimensions(self): + def cube_like_dimensions(self) -> tuple[u.Quantity, ...]: """ The length of each array axis as if all cubes were concatenated along the common axis. """ @@ -85,7 +89,7 @@ def cube_like_dimensions(self): return tuple(u.Quantity(d, unit=u.pix) for d in self.cube_like_shape) @property - def cube_like_shape(self): + def cube_like_shape(self) -> list[Any]: """ The length of each array axis as if all cubes were concatenated along the common axis. """ @@ -100,17 +104,17 @@ def cube_like_shape(self): return cube_like_shape @property - def cube_like_array_axis_physical_types(self): + def cube_like_array_axis_physical_types(self) -> list[tuple[str, ...]]: """ The physical types associated with each array axis, omitting the sequence axis. """ if self._common_axis is None: raise ValueError("Common axis must be set.") - return self.data[0].array_axis_physical_types + return self.data[0].array_axis_physical_types # type: ignore[no-any-return] - def __getitem__(self, item): + def __getitem__(self, item: Any) -> Any: if isinstance(item, numbers.Integral): - return self.data[item] + return self.data[int(item)] # Determine whether meta attribute should be sliced. new_meta = self.meta.slice[item] if (hasattr(self.meta, "__ndcube_can_slice__") and self.meta.__ndcube_can_slice__) else copy.deepcopy(self.meta) # Create an empty sequence in which to place the sliced cubes. @@ -119,7 +123,7 @@ def __getitem__(self, item): result.data = self.data[item] else: if isinstance(item[0], numbers.Integral): - result = self.data[item[0]][item[1:]] + result = self.data[int(item[0])][item[1:]] else: result.data = [cube[item[1:]] for cube in self.data[item[0]]] # Determine common axis after slicing. @@ -134,7 +138,7 @@ def __getitem__(self, item): return result @property - def index_as_cube(self): + def index_as_cube(self) -> "_IndexAsCubeSlicer": """ Slice the NDCubesequence instance as a single cube concatenated along the common axis. @@ -153,7 +157,7 @@ def index_as_cube(self): return _IndexAsCubeSlicer(self) @property - def common_axis_coords(self): + def common_axis_coords(self) -> list[list[Any]]: """ The coordinate values at each location along the common axis across all cubes. @@ -170,7 +174,7 @@ def common_axis_coords(self): cube_wcs = cube.combined_wcs common_coords.append(cube.axis_world_coords(common_axis, wcs=cube_wcs)) mappings.append(utils.wcs.array_indices_for_world_objects(cube_wcs, - axes=(common_axis,))) + axes=(common_axis,))) # type: ignore[arg-type] # For each coordinate, break up and then combine the coordinate objects across # the cubes into a list of coordinate objects that are length-1 and sequential # along the common axis. @@ -180,7 +184,7 @@ def common_axis_coords(self): for cube_idx in range(len(common_coords)): coord = common_coords[cube_idx][coord_idx] axis = np.where(np.array(mappings[cube_idx][coord_idx]) == common_axis)[0][0] - item = [slice(None)] * len(coord.shape) + item: list[Any] = [slice(None)] * len(coord.shape) for i in range(coord.shape[axis]): item[axis] = i exploded_coord.append(coord[tuple(item)]) @@ -188,7 +192,7 @@ def common_axis_coords(self): return sequence_coords @property - def sequence_axis_coords(self): + def sequence_axis_coords(self) -> dict[Any, list[Any]]: """ Return the coordinate values along the sequence axis. @@ -202,7 +206,7 @@ def sequence_axis_coords(self): return {name: [cube.global_coords[name] for cube in self.data] for name in global_names} - def explode_along_axis(self, axis): + def explode_along_axis(self, axis: int) -> "NDCubeSequenceBase": """ Separates slices of N-D cubes along a given cube axis into (N-1)D cubes. @@ -222,7 +226,7 @@ def explode_along_axis(self, axis): # To store the resultant cube result_cubes = [] # All slices are initially initialised as slice(None, None, None) - result_cubes_slice = [slice(None, None, None)] * len(self[0].data.shape) + result_cubes_slice: list[Any] = [slice(None, None, None)] * len(self[0].data.shape) # the range of the axis that needs to be sliced range_of_axis = self[0].data.shape[axis] for ndcube in self.data: @@ -236,12 +240,12 @@ def explode_along_axis(self, axis): new_common_axis = None elif self._common_axis > axis: new_common_axis = self._common_axis - 1 - elif self._common_axis < axis: + else: new_common_axis = self._common_axis # creating a new sequence with the result_cubes keeping the meta and common axis as axis return self._new_instance(result_cubes, common_axis=new_common_axis, meta=self.meta) - def crop(self, *points, wcses=None): + def crop(self, *points: Any, wcses: Any = None) -> Any: """ Crop cubes in sequence to smallest pixel-space bounding box containing the input points. @@ -277,7 +281,7 @@ def crop(self, *points, wcses=None): item = self._get_sequence_crop_item(*points, wcses=wcses) return self[item] - def crop_by_values(self, *points, units=None, wcses=None): + def crop_by_values(self, *points: Any, units: Any = None, wcses: Any = None) -> Any: """ Crop cubes in sequence to smallest pixel-space bounding box containing the input points. @@ -315,7 +319,8 @@ def crop_by_values(self, *points, units=None, wcses=None): item = self._get_sequence_crop_item(*points, wcses=wcses, crop_by_values=True, units=units) return self[item] - def _get_sequence_crop_item(self, *points, wcses=None, crop_by_values=False, units=None): + def _get_sequence_crop_item(self, *points: Any, wcses: Any = None, crop_by_values: bool = False, + units: Any = None) -> tuple[slice, ...]: """ Get the slice item with which to crop an NDCubeSequence given crop inputs. @@ -378,7 +383,7 @@ def _get_sequence_crop_item(self, *points, wcses=None, crop_by_values=False, uni return tuple( [slice(0, n_cubes)] + [slice(start, stop) for start, stop in zip(starts, stops)]) - def __str__(self): + def __str__(self) -> str: return (textwrap.dedent(f"""\ NDCubeSequence -------------- @@ -386,17 +391,18 @@ def __str__(self): Physical Types of Axes: {self.array_axis_physical_types} Common Cube Axis: {self._common_axis}""")) - def __repr__(self): + def __repr__(self) -> str: return f"{object.__repr__(self)}\n{self!s}" - def __len__(self): + def __len__(self) -> int: return len(self.data) - def __iter__(self): + def __iter__(self) -> Iterator[Any]: return iter(self.data) @classmethod - def _new_instance(cls, data_list, meta=None, common_axis=None): + def _new_instance(cls, data_list: Sequence[Any], meta: Any = None, + common_axis: int | None = None) -> "NDCubeSequenceBase": """ Instantiate a new instance of this class using given data. """ @@ -436,7 +442,7 @@ class NDCubeSequence(NDCubeSequenceBase): # last moment. plotter = PlotterDescriptor(default_type="mpl_sequence_plotter") - def plot(self, *args, **kwargs): + def plot(self, *args: Any, **kwargs: Any) -> Any: """ A convenience function for the plotters default ``plot()`` method. @@ -452,7 +458,7 @@ def plot(self, *args, **kwargs): return self.plotter.plot(*args, **kwargs) - def plot_as_cube(self, *args, **kwargs): + def plot_as_cube(self, *args: Any, **kwargs: Any) -> None: raise NotImplementedError( "NDCubeSequence plot_as_cube is no longer supported.\n" "To learn why or to tell us why it should be re-instated, " @@ -478,11 +484,15 @@ class _IndexAsCubeSlicer: Object of NDCubeSequence. """ - def __init__(self, seq): + def __init__(self, seq: NDCubeSequenceBase) -> None: self.seq = seq - def __getitem__(self, item): - common_axis = self.seq._common_axis + def __getitem__(self, item: Any) -> Any: + # _IndexAsCubeSlicer is only constructed by NDCubeSequenceBase.index_as_cube, + # which already guarantees _common_axis is not None. + if self.seq._common_axis is None: # pyright: ignore[reportPrivateUsage] + raise RuntimeError("common_axis must not be None here.") + common_axis = self.seq._common_axis # pyright: ignore[reportPrivateUsage] common_axis_lengths = [int(cube.shape[common_axis]) for cube in self.seq.data] n_cube_dims = len(self.seq.cube_like_shape) n_uncommon_cube_dims = n_cube_dims - 1 diff --git a/ndcube/tests/test_ndcollection.py b/ndcube/tests/test_ndcollection.py index 158112cd9..a604f66c6 100644 --- a/ndcube/tests/test_ndcollection.py +++ b/ndcube/tests/test_ndcollection.py @@ -112,13 +112,45 @@ def test_collection_pop(collection, popped_key, expected_popped, expected_collec @pytest.mark.parametrize(('collection', 'key', 'expected'), [ (cube_collection, "cube0", NDCollection([("cube1", cube1), ("cube2", cube2)], - aligned_axes=aligned_axes[1:]))]) + aligned_axes=aligned_axes[1:])), + # Regression test: __delitem__ used to call self.aligned_axes.__delitem__(key) + # unconditionally, crashing with AttributeError on a collection with no + # aligned axes since self.aligned_axes is None in that case. + (unaligned_collection, "cube0", NDCollection([("cube1", cube1), ("cube2", cube2)]))]) def test_del_collection(collection, key, expected): del_collection = collection.copy() del del_collection[key] helpers.assert_collections_equal(del_collection, expected) +def test_slice_by_keys_unaligned_collection(): + # Regression test: slicing by a sequence of string keys used to crash with + # TypeError ("'NoneType' object is not subscriptable") on a collection with + # aligned_axes=None, because `self.aligned_axes[item_]` was called + # unconditionally instead of being guarded by an is-None check. + result = unaligned_collection[("cube0", "cube1")] + assert list(result.keys()) == ["cube0", "cube1"] + assert result.aligned_axes is None + + +def test_collection_getitem_unsupported_type(): + # Regression test: the "unsupported slicing type" error message referenced + # `axis_item`, a loop variable only bound inside the tuple-handling branch, + # so hitting this branch raised NameError instead of the intended TypeError. + with pytest.raises(TypeError, match="Unsupported slicing type"): + cube_collection[1.5] + + +def test_collection_slice_open_start_negative_stop_with_sliceable_meta(): + # Regression test: sanitizing a slice item with sliceable meta compared + # `item.start < 0` even when `item.start` is None (i.e. an open-ended + # slice like `slice(None, -1)`), crashing with + # "'<' not supported between instances of 'NoneType' and 'int'". + result = cube_collection[slice(None, -1)] + expected = cube_collection[slice(0, cube_collection.aligned_dimensions[0] - 1)] + helpers.assert_collections_equal(result, expected) + + @pytest.mark.parametrize(('collection', 'key', 'data', 'aligned_axes', 'expected'), [ (cube_collection, "cube1", cube2, aligned_axes[2], NDCollection( [("cube0", cube0), ("cube1", cube2), ("cube2", cube2)], diff --git a/ndcube/tests/test_ndcube_slice_and_crop.py b/ndcube/tests/test_ndcube_slice_and_crop.py index 56b7cd4a4..48d5711f4 100644 --- a/ndcube/tests/test_ndcube_slice_and_crop.py +++ b/ndcube/tests/test_ndcube_slice_and_crop.py @@ -423,6 +423,17 @@ def test_crop_by_values_1d_dependent(ndcube_4d_ln_lt_l_t): helpers.assert_cubes_equal(output, expected) +def test_crop_no_points_with_extra_coords_wcs(ndcube_3d_ln_lt_l_ec_time): + # Regression test: cropping with zero points is a no-op, but + # sanitize_crop_inputs used to return the *un-unwrapped* ExtraCoords + # object on that early-exit path, so `wcs.pixel_n_dim` crashed with + # AttributeError since ExtraCoords has no such attribute (only + # `ExtraCoords.wcs.pixel_n_dim` does). + cube = ndcube_3d_ln_lt_l_ec_time + output = cube.crop(wcs=cube.extra_coords) + helpers.assert_cubes_equal(output, cube) + + def test_crop_by_extra_coords(ndcube_3d_ln_lt_l_ec_time): cube = ndcube_3d_ln_lt_l_ec_time lower_corner = (Time("2000-01-01T15:00:00", scale="utc", format="fits"), None) diff --git a/ndcube/utils/collection.py b/ndcube/utils/collection.py index c9a84e6e8..b14abf020 100644 --- a/ndcube/utils/collection.py +++ b/ndcube/utils/collection.py @@ -1,11 +1,15 @@ import numbers +from typing import Any +from collections.abc import Sequence import numpy as np +import numpy.typing as npt __all__ = ['assert_aligned_axes_compatible'] -def _sanitize_aligned_axes(keys, data, aligned_axes): +def _sanitize_aligned_axes(keys: Sequence[Any], data: Sequence[Any], # pyright: ignore[reportUnusedFunction] + aligned_axes: str | tuple[Any, ...] | int | None) -> dict[Any, tuple[Any, ...]] | None: if aligned_axes is None: return None # If aligned_axes set to "all", assume all axes are aligned in order. @@ -25,7 +29,7 @@ def _sanitize_aligned_axes(keys, data, aligned_axes): return dict(zip(keys, sanitized_axes)) -def _sanitize_user_aligned_axes(data, aligned_axes): +def _sanitize_user_aligned_axes(data: Sequence[Any], aligned_axes: Any) -> tuple[Any, ...]: """ Converts input aligned_axes to standard format. aligned_axes can be supplied by the user in a few ways: @@ -55,7 +59,7 @@ def _sanitize_user_aligned_axes(data, aligned_axes): n_cubes = len(data) if axes_all_ints: n_aligned_axes = len(aligned_axes) - aligned_axes = tuple([aligned_axes for i in range(n_cubes)]) + aligned_axes = tuple([aligned_axes for _ in range(n_cubes)]) # If all elements are tuple, ensure there is a tuple for each cube and # all elements of each sub-tuple are ints. @@ -110,30 +114,30 @@ def _sanitize_user_aligned_axes(data, aligned_axes): return aligned_axes -def _update_aligned_axes(drop_aligned_axes_indices, aligned_axes, first_key): +def _update_aligned_axes(drop_aligned_axes_indices: npt.NDArray[Any], aligned_axes: dict[Any, Any], # pyright: ignore[reportUnusedFunction] + first_key: Any) -> tuple[Any, ...] | None: # Remove dropped axes from aligned_axes. MUST BE A BETTER WAY TO DO THIS. if len(drop_aligned_axes_indices) <= 0: - new_aligned_axes = tuple(aligned_axes.values()) - elif len(drop_aligned_axes_indices) == len(aligned_axes[first_key]): - new_aligned_axes = None - else: - new_aligned_axes = [] - for key in aligned_axes.keys(): - cube_aligned_axes = np.array(aligned_axes[key]) - for drop_axis_index in drop_aligned_axes_indices: - drop_axis = cube_aligned_axes[drop_axis_index] - cube_aligned_axes = np.delete(cube_aligned_axes, drop_axis_index) - w = np.where(cube_aligned_axes > drop_axis) - cube_aligned_axes[w] -= 1 - w = np.where(drop_aligned_axes_indices > drop_axis_index) - drop_aligned_axes_indices[w] -= 1 - new_aligned_axes.append(tuple(cube_aligned_axes)) - new_aligned_axes = tuple(new_aligned_axes) - - return new_aligned_axes - - -def assert_aligned_axes_compatible(data_dimensions1, data_dimensions2, data_axes1, data_axes2): + return tuple(aligned_axes.values()) + if len(drop_aligned_axes_indices) == len(aligned_axes[first_key]): + return None + + new_aligned_axes: list[Any] = [] + for axes in aligned_axes.values(): + cube_aligned_axes = np.array(axes) + for drop_axis_index in drop_aligned_axes_indices: + drop_axis = cube_aligned_axes[drop_axis_index] + cube_aligned_axes = np.delete(cube_aligned_axes, drop_axis_index) + w = np.where(cube_aligned_axes > drop_axis) + cube_aligned_axes[w] -= 1 + w = np.where(drop_aligned_axes_indices > drop_axis_index) + drop_aligned_axes_indices[w] -= 1 + new_aligned_axes.append(tuple(cube_aligned_axes)) + return tuple(new_aligned_axes) + + +def assert_aligned_axes_compatible(data_dimensions1: tuple[int, ...], data_dimensions2: tuple[int, ...], + data_axes1: tuple[int, ...] | None, data_axes2: tuple[int, ...] | None) -> None: """ Checks whether two sets of aligned axes are compatible. @@ -143,10 +147,12 @@ def assert_aligned_axes_compatible(data_dimensions1, data_dimensions2, data_axes The dimension lengths of data cube 1. data_dimensions2: `tuple` of `int` The dimension lengths of data cube 2. - data_axes1: `tuple` of `int` - The aligned axes of data cube 1. - data_axes2: `tuple` of `int` - The aligned axes of data cube 2. + data_axes1: `tuple` of `int` or `None` + The aligned axes of data cube 1. `None` if cube 1's collection has no + aligned axes. + data_axes2: `tuple` of `int` or `None` + The aligned axes of data cube 2. `None` if cube 2's collection has no + aligned axes. """ # If one set of aligned axes is None and the other isn't, they are defined as not compatible. if (data_axes1 is None and data_axes2 is not None) or (data_axes1 is not None and data_axes2 is None): @@ -155,6 +161,9 @@ def assert_aligned_axes_compatible(data_dimensions1, data_dimensions2, data_axes # Aligned_axes are being used for both collections if data_axes1 is not None: + # Guaranteed by the check above: if data_axes1 is not None, data_axes2 isn't either. + if data_axes2 is None: + raise RuntimeError("data_axes2 must not be None here.") # Confirm same number of aligned axes. if len(data_axes1) != len(data_axes2): raise ValueError(f"Number of aligned axes must be equal: {len(data_axes1)} != {len(data_axes2)}") diff --git a/ndcube/utils/cube.py b/ndcube/utils/cube.py index d1c188a13..82a748959 100644 --- a/ndcube/utils/cube.py +++ b/ndcube/utils/cube.py @@ -1,8 +1,11 @@ import inspect +from typing import TYPE_CHECKING, Any, TypeAlias from functools import wraps from itertools import chain +from collections.abc import Callable, Iterable import numpy as np +import numpy.typing as npt import astropy.nddata from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS, HighLevelWCSWrapper, SlicedLowLevelWCS @@ -10,6 +13,11 @@ from ndcube.utils import wcs as wcs_utils from ndcube.utils.exceptions import warn_user +if TYPE_CHECKING: + from ndcube.extra_coords.extra_coords import ExtraCoordsABC + +WCSType: TypeAlias = BaseHighLevelWCS | BaseLowLevelWCS + __all__ = [ "get_crop_item_from_points", "propagate_rebin_uncertainties", @@ -18,7 +26,7 @@ ] -def sanitize_wcs(func): +def sanitize_wcs(func: Callable[..., Any]) -> Callable[..., Any]: """ A wrapper for NDCube methods to sanitise the wcs argument. @@ -34,7 +42,7 @@ def sanitize_wcs(func): from ndcube.extra_coords.extra_coords import ExtraCoords # noqa: PLC0415 @wraps(func) - def wcs_wrapper(*args, **kwargs): + def wcs_wrapper(*args: Any, **kwargs: Any) -> Any: sig = inspect.signature(func) params = sig.bind(*args, **kwargs) wcs = params.arguments.get('wcs', None) @@ -60,7 +68,8 @@ def wcs_wrapper(*args, **kwargs): return wcs_wrapper -def sanitize_crop_inputs(points, wcs): +def sanitize_crop_inputs(points: Iterable[Any], wcs: 'WCSType | ExtraCoordsABC' + ) -> tuple[bool, list[Any], WCSType]: """Sanitize inputs to NDCube crop methods. First arg returned signifies whether the inputs imply that cropping @@ -68,7 +77,7 @@ def sanitize_crop_inputs(points, wcs): """ points = list(points) n_points = len(points) - n_coords = [None] * n_points + n_coords: list[int | None] = [None] * n_points values_are_none = [False] * n_points for i, point in enumerate(points): # Ensure each point is a list @@ -84,30 +93,31 @@ def sanitize_crop_inputs(points, wcs): values_are_none[i] = True # Squeeze length-1 coordinate objects to scalars. points[i] = [coord.squeeze() if hasattr(coord, "squeeze") else coord for coord in points[i]] - # If no points contain a coord, i.e. if all entries in all points are None, - # set no-op flag to True and exit. - if all(values_are_none): - return True, points, wcs - # Not not all points are of same length, error. - if len(set(n_coords)) != 1: - raise ValueError("All points must have same number of coordinate objects." - f"Number of objects in each point: {n_coords}") # Import must be here to avoid circular import. from ndcube.extra_coords.extra_coords import ExtraCoords # noqa: PLC0415 if isinstance(wcs, ExtraCoords): # Determine how many dummy axes are needed - n_dummy_axes = len(wcs._cube_array_axes_without_extra_coords) + n_dummy_axes = len(wcs._cube_array_axes_without_extra_coords) # pyright: ignore[reportPrivateUsage] if n_dummy_axes > 0: points = [point + [None] * n_dummy_axes for point in points] # Convert extra coords to WCS describing whole cube. wcs = wcs.cube_wcs # Ensure WCS is low level. if isinstance(wcs, BaseHighLevelWCS): - wcs = wcs.low_level_wcs + wcs = wcs.low_level_wcs # pyright: ignore[reportAttributeAccessIssue] + # If no points contain a coord, i.e. if all entries in all points are None, + # set no-op flag to True and exit. + if all(values_are_none): + return True, points, wcs + # Not not all points are of same length, error. + if len(set(n_coords)) != 1: + raise ValueError("All points must have same number of coordinate objects." + f"Number of objects in each point: {n_coords}") return False, points, wcs -def get_crop_item_from_points(points, wcs, crop_by_values, keepdims, original_shape): +def get_crop_item_from_points(points: Iterable[Iterable[Any]], wcs: WCSType, crop_by_values: bool, + keepdims: bool, original_shape: tuple[int, ...]) -> tuple[int | slice, ...]: """ Find slice item that crops to minimum cube in array-space containing specified world points. @@ -136,15 +146,16 @@ def get_crop_item_from_points(points, wcs, crop_by_values, keepdims, original_sh Returns ------- - item : `tuple` of `slice` + item : `tuple` of `int` or ``slice`` The slice item for each axis of the cube which, when applied to the cube, will return the minimum cube in array-index-space that contains all the - input world points. + input world points. An axis collapsed to a single element (when + ``keepdims=False``) is given as an `int` rather than a length-1 ``slice``. """ # Define a list of lists to hold the pixel coordinates of the points # where each inner list gives the pixel coordinates of all points for that pixel axis. # Recall that pixel axis ordering is reversed compared to array axis ordering. - combined_points_pixel_idx = [[]] * wcs.pixel_n_dim + combined_points_pixel_idx: list[list[Any]] = [[]] * wcs.pixel_n_dim high_level_wcs = HighLevelWCSWrapper(wcs) if isinstance(wcs, BaseLowLevelWCS) else wcs low_level_wcs = high_level_wcs.low_level_wcs # For each point compute the corresponding array indices. @@ -206,7 +217,7 @@ def get_crop_item_from_points(points, wcs, crop_by_values, keepdims, original_sh # and then convert to array indices. Note that combined_points_pixel_idx holds the # pixel coords for each pixel axis. Therefore, to iterate in array axis order, # combined_points_pixel_idx must be reversed. - item = [] + item: list[int | slice] = [] ambiguous = False message = "" result_is_scalar = True @@ -270,8 +281,11 @@ def get_crop_item_from_points(points, wcs, crop_by_values, keepdims, original_sh return tuple(item) -def propagate_rebin_uncertainties(uncertainty, data, mask, operation, operation_ignores_mask=False, - propagation_operation=None, correlation=0, **kwargs): +def propagate_rebin_uncertainties(uncertainty: astropy.nddata.NDUncertainty, data: npt.NDArray[Any], + mask: npt.NDArray[np.bool_] | bool | None, operation: Callable[..., Any], + operation_ignores_mask: bool = False, + propagation_operation: Callable[..., Any] | None = None, + correlation: int = 0, **kwargs: Any) -> astropy.nddata.NDUncertainty: """ Default algorithm for uncertainty propagation in :meth:`~ndcube.NDCube.rebin`. @@ -287,7 +301,7 @@ def propagate_rebin_uncertainties(uncertainty, data, mask, operation, operation_ The uncertainties associated with the data. The first dimension represents pixels in each bin being aggregated while trailing dimensions must have the same shape as the rebinned data. - data: array-like or `None` + data: array-like The data associated with the above uncertainties. Must have same shape as above. mask: array-like of `bool` or `None` @@ -328,40 +342,47 @@ def propagate_rebin_uncertainties(uncertainty, data, mask, operation, operation_ raise ValueError("propagation_operation not recognized.") # Build mask if not provided. new_uncertainty = uncertainty[0] # Define uncertainty for initial iteration step. + # idx and n_pix_per_bin are each only read below in the branch correlated + # with the code path that sets so the None defaults are never read. + idx: npt.NDArray[np.bool_] | None = None + n_pix_per_bin: int | None = None if operation_ignores_mask or mask is None: mask = False if mask is False: - if operation_is_nantype: - nan_mask = np.isnan(data) - if nan_mask.any(): - mask = nan_mask - idx = np.logical_not(mask) - mask1 = mask[1:] + nan_mask = np.isnan(data) if operation_is_nantype else None + if nan_mask is not None and nan_mask.any(): + mask = nan_mask + idx = np.logical_not(mask) + mask1 = mask[1:] else: - # If there is no mask and operation is not nan-type, build generator - # so non-mask can still be iterated. + # If there is no mask and operation is not nan-type (or no NaNs are + # actually present), build a generator so non-mask can still be iterated. n_pix_per_bin = data.shape[flat_axis] - mask1 = (False for i in range(1, n_pix_per_bin)) + # n_pix_per_bin is provably an int here, but pyright can't narrow + # it through the generator's closure. + mask1 = (False for _ in range(1, n_pix_per_bin)) # pyright: ignore[reportArgumentType] else: # Mask uncertainties corresponding to nan data if operation is nantype. + # mask is guaranteed to be an array here (not the `False` sentinel above), + # since callers only pass array-like masks or None per the docstring. if operation_is_nantype: - mask[np.isnan(data)] = True + mask[np.isnan(data)] = True # type: ignore[index] # Set masked uncertainties in first mask to 0 # as they shouldn't count towards final uncertainty. - mask1 = mask[1:] + mask1 = mask[1:] # type: ignore[index] idx = np.logical_not(mask) uncertainty.array[mask] = 0 - new_uncertainty.array[mask[0]] = 0 + new_uncertainty.array[mask[0]] = 0 # type: ignore[index] # Propagate uncertainties. # Note uncertainty must be associated with a parent nddata for some propagations. cumul_data = data[0] if mask is not False and operation_ignores_mask is False: - cumul_data[idx[0]] = 0 + cumul_data[idx[0]] = 0 # type: ignore[index] parent_nddata = astropy.nddata.NDData(cumul_data, uncertainty=new_uncertainty) new_uncertainty.parent_nddata = parent_nddata for j, mask_slice in enumerate(mask1): i = j + 1 - cumul_data = operation(data[:i+1]) if mask is False else operation(data[:i+1][idx[:i+1]]) + cumul_data = operation(data[:i+1]) if mask is False else operation(data[:i+1][idx[:i+1]]) # type: ignore[index] data_slice = astropy.nddata.NDData(data=data[i], mask=mask_slice, uncertainty=uncertainty[i]) new_uncertainty = new_uncertainty.propagate(propagation_operation, data_slice, diff --git a/ndcube/utils/exceptions.py b/ndcube/utils/exceptions.py index ce98fa833..3ad47144c 100644 --- a/ndcube/utils/exceptions.py +++ b/ndcube/utils/exceptions.py @@ -29,7 +29,7 @@ class NDCubeDeprecationWarning(FutureWarning, NDCubeWarning): """ -def warn_user(msg, stacklevel=1): +def warn_user(msg: str, stacklevel: int = 1) -> None: """ Raise a `NDCubeWarning`. @@ -45,7 +45,7 @@ def warn_user(msg, stacklevel=1): warnings.warn(msg, NDCubeUserWarning , stacklevel + 1) -def warn_deprecated(msg, stacklevel=1): +def warn_deprecated(msg: str, stacklevel: int = 1) -> None: """ Raise a `NDCubeDeprecationWarning`. diff --git a/ndcube/utils/misc.py b/ndcube/utils/misc.py index 2dee74c69..1f30fd24f 100644 --- a/ndcube/utils/misc.py +++ b/ndcube/utils/misc.py @@ -1,17 +1,20 @@ +from typing import Any +from collections.abc import Iterable + import astropy.units as u __all__ = ['convert_quantities_to_units', 'unique_sorted'] -def unique_sorted(iterable): +def unique_sorted(iterable: Iterable[Any]) -> list[Any]: """ Return unique values in the order they are first encountered in the iterable. """ - lookup = set() # a temporary lookup set - return [ele for ele in iterable if ele not in lookup and lookup.add(ele) is None] + lookup: set[Any] = set() # a temporary lookup set + return [ele for ele in iterable if ele not in lookup and lookup.add(ele) is None] # type: ignore[func-returns-value] -def convert_quantities_to_units(coords, units): +def convert_quantities_to_units(coords: Iterable[Any], units: Iterable[Any]) -> list[Any]: """Converts a sequence of Quantities to units used in the WCS. Non-Quantity types in the sequence are allowed and ignored. diff --git a/ndcube/utils/sequence.py b/ndcube/utils/sequence.py index 05af52585..cc4de469e 100644 --- a/ndcube/utils/sequence.py +++ b/ndcube/utils/sequence.py @@ -4,7 +4,8 @@ """ from copy import deepcopy -from collections import namedtuple +from typing import Any, NamedTuple +from collections.abc import Sequence import numpy as np @@ -13,17 +14,18 @@ 'cube_like_tuple_item_to_sequence_items'] -SequenceItem = namedtuple("SequenceItem", "sequence_index cube_item") -""" -Define SequenceItem named tuple of length 2. Its attributes are: -sequence_index: an int giving the index of a cube within an NDCubeSequence. -cube_item: item (int, slice, tuple) to be applied to cube identified -by sequence_index attribute. -""" +class SequenceItem(NamedTuple): + """ + sequence_index: an int giving the index of a cube within an NDCubeSequence. + cube_item: item (int, slice, tuple) to be applied to cube identified + by sequence_index attribute. + """ + sequence_index: int + cube_item: Any -def cube_like_index_to_sequence_and_common_axis_indices(cube_like_index, common_axis, - common_axis_lengths): +def cube_like_index_to_sequence_and_common_axis_indices(cube_like_index: int, common_axis: int, + common_axis_lengths: Sequence[int]) -> tuple[int, int]: """ Converts a cube-like index for an NDCubeSequence to a sequence index and a common axis index. @@ -52,10 +54,12 @@ def cube_like_index_to_sequence_and_common_axis_indices(cube_like_index, common_ common_axis_index = cube_like_index else: common_axis_index = cube_like_index - cumul_lengths[sequence_index - 1] - return sequence_index, common_axis_index + return int(sequence_index), int(common_axis_index) -def cube_like_tuple_item_to_sequence_items(item, common_axis, common_axis_lengths, n_cube_dims): +def cube_like_tuple_item_to_sequence_items(item: list[Any], common_axis: int, + common_axis_lengths: Sequence[int], + n_cube_dims: int) -> list[SequenceItem]: """ Convert a tuple for slicing an NDCubeSequence in the cube-like API to a list of SequenceItems. @@ -64,8 +68,8 @@ def cube_like_tuple_item_to_sequence_items(item, common_axis, common_axis_length Parameters ---------- - item: iterable of `int` or `slice` - The slicing item. The common axis entry must be a `slice` + item: iterable of `int` or ``slice`` + The slicing item. The common axis entry must be a ``slice`` common_axis: `int` The index of the item corresponding to the common axis. diff --git a/ndcube/utils/sphinx/code_context.py b/ndcube/utils/sphinx/code_context.py index 1048dcdd0..177cd3e54 100644 --- a/ndcube/utils/sphinx/code_context.py +++ b/ndcube/utils/sphinx/code_context.py @@ -3,6 +3,8 @@ user a way to see that context. """ +from typing import Any + from docutils import nodes from docutils.parsers.rst import directives from sphinx.directives.code import CodeBlock @@ -31,7 +33,7 @@ class ExpandingCodeBlock(CodeBlock): 'summary': directives.unchanged_required, } - def run(self): + def run(self) -> list[Any]: source, lineno = self.state_machine.get_source_and_line(self.lineno) summary_text = self.options.get('summary', 'Show setup code') @@ -51,7 +53,7 @@ def run(self): return [open_raw_node, literal, close_raw_node] -def setup(app): +def setup(app: Any) -> dict[str, bool]: app.add_directive('expanding-code-block', ExpandingCodeBlock) return {'parallel_read_safe': True, 'parallel_write_safe': True} diff --git a/ndcube/utils/tests/test_utils_cube.py b/ndcube/utils/tests/test_utils_cube.py index 5f410e435..456e1fd44 100644 --- a/ndcube/utils/tests/test_utils_cube.py +++ b/ndcube/utils/tests/test_utils_cube.py @@ -4,7 +4,7 @@ from astropy.nddata import StdDevUncertainty -from ndcube.utils.cube import propagate_rebin_uncertainties +from ndcube.utils.cube import propagate_rebin_uncertainties, sanitize_crop_inputs @pytest.fixture @@ -99,3 +99,18 @@ def test_propagate_rebin_uncertainties_nan(stacked_pixel_data): np.nanmean, operation_ignores_mask=False) assert type(output) is type(expected) assert np.allclose(output.array, expected.array) + + +def test_sanitize_crop_inputs_no_op_unwraps_extra_coords(ndcube_3d_ln_lt_l_ec_time): + # Regression test: on the no-op (all-points-None, including zero points) + # early-exit path, sanitize_crop_inputs used to return the wcs argument + # completely unprocessed, so if it was passed an ExtraCoords object it + # would still be an ExtraCoords object rather than being unwrapped to a + # low-level WCS like on the normal path. Callers (e.g. NDCube._get_crop_item) + # immediately use low-level WCS attributes like `.pixel_n_dim` on the + # returned value, which ExtraCoords does not have. + cube = ndcube_3d_ln_lt_l_ec_time + no_op, points, wcs = sanitize_crop_inputs((), cube.extra_coords) + assert no_op is True + assert points == [] + assert hasattr(wcs, "pixel_n_dim") diff --git a/ndcube/utils/wcs.py b/ndcube/utils/wcs.py index 244a59316..4ef63a9dc 100644 --- a/ndcube/utils/wcs.py +++ b/ndcube/utils/wcs.py @@ -3,13 +3,20 @@ """ import numbers +from typing import Any, cast from collections import UserDict +from collections.abc import Iterable, Sequence import numpy as np +import numpy.typing as npt from astropy.wcs.utils import pixel_to_pixel from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS, low_level_api +# A pixel/world axis index, as accepted and returned throughout this module. +# Indices returned by numpy indexing are numpy integers, not plain ints. +AxisIndex = int | np.integer[Any] + __all__ = [ 'array_indices_for_world_objects', 'calculate_world_indices_from_axes', @@ -29,9 +36,9 @@ ] -class TwoWayDict(UserDict): +class TwoWayDict(UserDict[str, str]): @property - def inv(self): + def inv(self) -> dict[str, str]: """ The inverse dictionary. """ @@ -62,7 +69,7 @@ def inv(self): wcs_ivoa_mapping[key] = value -def convert_between_array_and_pixel_axes(axis, naxes): +def convert_between_array_and_pixel_axes(axis: npt.NDArray[np.integer[Any]], naxes: int) -> npt.NDArray[np.integer[Any]]: """Reflects axis index about center of number of axes. This is used to convert between array axes in numpy order and pixel axes in WCS order. @@ -82,7 +89,7 @@ def convert_between_array_and_pixel_axes(axis, naxes): The axis number(s) after reflection. """ # Check type of input. - if not isinstance(axis, np.ndarray): + if not isinstance(axis, np.ndarray): # pyright: ignore[reportUnnecessaryIsInstance] raise TypeError(f"input must be of array type. Got type: {type(axis)}") if axis.dtype.char not in np.typecodes['AllInteger']: raise TypeError(f"input dtype must be of int type. Got dtype: {axis.dtype})") @@ -96,7 +103,7 @@ def convert_between_array_and_pixel_axes(axis, naxes): -def pixel_axis_to_world_axes(pixel_axis, axis_correlation_matrix): +def pixel_axis_to_world_axes(pixel_axis: AxisIndex, axis_correlation_matrix: npt.NDArray[np.bool_]) -> npt.NDArray[np.integer[Any]]: """ Retrieves the indices of the world axis physical types corresponding to a pixel axis. @@ -117,7 +124,7 @@ def pixel_axis_to_world_axes(pixel_axis, axis_correlation_matrix): return np.arange(axis_correlation_matrix.shape[0])[axis_correlation_matrix[:, pixel_axis]] -def world_axis_to_pixel_axes(world_axis, axis_correlation_matrix): +def world_axis_to_pixel_axes(world_axis: AxisIndex, axis_correlation_matrix: npt.NDArray[np.bool_]) -> npt.NDArray[np.integer[Any]]: """ Gets the pixel axis indices corresponding to the index of a world axis. @@ -135,10 +142,13 @@ def world_axis_to_pixel_axes(world_axis, axis_correlation_matrix): pixel_axes: `numpy.ndarray` The pixel axis indices corresponding to the world axis. """ - return np.arange(axis_correlation_matrix.shape[1])[axis_correlation_matrix[world_axis]] + return cast( + "npt.NDArray[np.integer[Any]]", + np.arange(axis_correlation_matrix.shape[1])[axis_correlation_matrix[world_axis]], + ) -def pixel_axis_to_physical_types(pixel_axis, wcs): +def pixel_axis_to_physical_types(pixel_axis: AxisIndex, wcs: BaseLowLevelWCS) -> npt.NDArray[np.str_]: """ Gets the world axis physical types corresponding to a pixel axis. @@ -155,17 +165,20 @@ def pixel_axis_to_physical_types(pixel_axis, wcs): physical_types: `numpy.ndarray` of `str` The physical types corresponding to the pixel axis. """ - return np.array(wcs.world_axis_physical_types)[wcs.axis_correlation_matrix[:, pixel_axis]] + return cast( + npt.NDArray[np.str_], + np.array(wcs.world_axis_physical_types)[wcs.axis_correlation_matrix[:, pixel_axis]], + ) -def physical_type_to_pixel_axes(physical_type, wcs): +def physical_type_to_pixel_axes(physical_type: str, wcs: BaseLowLevelWCS) -> npt.NDArray[np.integer[Any]]: """ Gets the pixel axis indices corresponding to a world axis physical type. Parameters ---------- - physical_type: `int` - The pixel axis number(s) for which the world axis numbers are desired. + physical_type: `str` + The physical type or a substring unique to a physical type. wcs: `astropy.wcs.wcsapi.BaseLowLevelWCS` The WCS object defining the relationship between pixel and world axes. @@ -179,7 +192,7 @@ def physical_type_to_pixel_axes(physical_type, wcs): return world_axis_to_pixel_axes(world_axis, wcs.axis_correlation_matrix) -def physical_type_to_world_axis(physical_type, world_axis_physical_types): +def physical_type_to_world_axis(physical_type: str, world_axis_physical_types: Sequence[str]) -> np.integer[Any]: """ Returns world axis index of a physical type based on WCS world_axis_physical_types. @@ -214,10 +227,10 @@ def physical_type_to_world_axis(physical_type, world_axis_physical_types): f" Got: {physical_type}" ) # Return axes with duplicates removed. - return widx[0] + return cast("np.integer[Any]", widx[0]) -def get_dependent_pixel_axes(pixel_axis, axis_correlation_matrix): +def get_dependent_pixel_axes(pixel_axis: AxisIndex, axis_correlation_matrix: npt.NDArray[np.bool_]) -> npt.NDArray[np.integer[Any]]: """ Find indices of all pixel axes associated with the world axes linked to the input pixel axis. @@ -253,7 +266,7 @@ def get_dependent_pixel_axes(pixel_axis, axis_correlation_matrix): return np.sort(np.nonzero((world_dep & axis_correlation_matrix).any(axis=0))[0]) -def get_dependent_array_axes(array_axis, axis_correlation_matrix): +def get_dependent_array_axes(array_axis: AxisIndex, axis_correlation_matrix: npt.NDArray[np.bool_]) -> npt.NDArray[np.integer[Any]]: """ Find indices of all array axes associated with the world axes linked to the input array axis. @@ -289,7 +302,7 @@ def get_dependent_array_axes(array_axis, axis_correlation_matrix): return np.sort(dependent_array_axes) -def get_dependent_world_axes(world_axis, axis_correlation_matrix): +def get_dependent_world_axes(world_axis: AxisIndex, axis_correlation_matrix: npt.NDArray[np.bool_]) -> npt.NDArray[np.integer[Any]]: """ Given a WCS world axis index, return indices of dependent WCS world axes. @@ -318,7 +331,7 @@ def get_dependent_world_axes(world_axis, axis_correlation_matrix): return np.sort(np.nonzero((pixel_dep & axis_correlation_matrix).any(axis=1))[0]) -def get_dependent_physical_types(physical_type, wcs): +def get_dependent_physical_types(physical_type: str, wcs: BaseLowLevelWCS) -> npt.NDArray[np.str_]: """ Given a world axis physical type, return the dependent physical types including the input type. @@ -338,10 +351,10 @@ def get_dependent_physical_types(physical_type, wcs): world_axis_physical_types = wcs.world_axis_physical_types world_axis = physical_type_to_world_axis(physical_type, world_axis_physical_types) dependent_world_axes = get_dependent_world_axes(world_axis, wcs.axis_correlation_matrix) - return np.array(world_axis_physical_types)[dependent_world_axes] + return cast("npt.NDArray[np.str_]", np.array(world_axis_physical_types)[dependent_world_axes]) -def validate_physical_types(physical_types): +def validate_physical_types(physical_types: Sequence[str | None]) -> None: """ Validate a list of physical types against the UCD1+ standard """ @@ -356,14 +369,14 @@ def validate_physical_types(physical_types): ) -def calculate_world_indices_from_axes(wcs, axes): +def calculate_world_indices_from_axes(wcs: BaseLowLevelWCS, axes: Iterable[AxisIndex | str]) -> npt.NDArray[np.integer[Any]]: """ Given a string representation of a world axis or a numerical array index, convert it to a numerical world index aligning to the position in wcs.world_axis_object_components. """ # Convert input axes to WCS world axis indices. - world_indices = [] + world_indices: list[np.integer[Any]] = [] for axis in axes: if isinstance(axis, numbers.Integral): # If axis is int, it is a numpy order array axis. @@ -383,7 +396,9 @@ def calculate_world_indices_from_axes(wcs, axes): return np.unique(np.array(world_indices, dtype=int)) -def pixel_indices_for_world_objects(wcs, axes=None): +def pixel_indices_for_world_objects( + wcs: BaseHighLevelWCS, axes: Iterable[AxisIndex | str] | None = None, +) -> tuple[npt.NDArray[np.integer[Any]], ...]: """ Calculate the pixel indices corresponding to each high level world object. @@ -405,12 +420,12 @@ def pixel_indices_for_world_objects(wcs, axes=None): Returns ------- - pixel_indices : `tuple` of `tuple` of `int` - For each world object, a tuple of pixel axes identified by their - number. Pixel indices in each sub-tuple are not guaranteed to be + pixel_indices : `tuple` of `numpy.ndarray` of `int` + For each world object, an array of pixel axes identified by their + number. Pixel indices in each sub-array are not guaranteed to be ordered with respect to the arrays in the object, as the object could be an object like ``SkyCoord`` where there is a separation of the two - coordinates. The pixel indices will be returned in the sub-tuple in + coordinates. The pixel indices will be returned in the sub-array in pixel index order. """ if axes: @@ -419,7 +434,7 @@ def pixel_indices_for_world_objects(wcs, axes=None): world_indices = np.arange(wcs.world_n_dim) object_names = np.array([wao_comp[0] for wao_comp in wcs.low_level_wcs.world_axis_object_components]) - pixel_indices = [None] * len(object_names) + pixel_indices: list[npt.NDArray[np.integer[Any]] | None] = [None] * len(object_names) for world_index, oname in enumerate(object_names): # If this world index is deselected by axes= then skip if world_index not in world_indices: @@ -434,7 +449,9 @@ def pixel_indices_for_world_objects(wcs, axes=None): return tuple(pi for pi in pixel_indices if pi is not None) -def array_indices_for_world_objects(wcs, axes=None): +def array_indices_for_world_objects( + wcs: BaseHighLevelWCS, axes: Iterable[AxisIndex | str] | None = None, +) -> tuple[tuple[np.integer[Any], ...], ...]: """ Calculate the array indices corresponding to each high level world object. @@ -465,14 +482,14 @@ def array_indices_for_world_objects(wcs, axes=None): array index order, i.e ascending. """ pixel_indices_of_world_objects = pixel_indices_for_world_objects(wcs, axes=axes) - array_indices_of_world_objects = [] + array_indices_of_world_objects: list[tuple[np.integer[Any], ...]] = [] for pixel_indices in pixel_indices_of_world_objects: array_indices = convert_between_array_and_pixel_axes(pixel_indices, wcs.pixel_n_dim) array_indices_of_world_objects.append(tuple(array_indices[::-1])) # Invert from pixel order to array order return tuple(array_indices_of_world_objects) -def get_low_level_wcs(wcs, name='wcs'): +def get_low_level_wcs(wcs: BaseHighLevelWCS | BaseLowLevelWCS, name: str = 'wcs') -> BaseLowLevelWCS: """ Returns a low level WCS object from a low level or high level WCS. @@ -496,7 +513,9 @@ def get_low_level_wcs(wcs, name='wcs'): raise ValueError(f'{name} must implement either BaseHighLevelWCS or BaseLowLevelWCS') -def compare_wcs_physical_types(source_wcs, target_wcs): +def compare_wcs_physical_types( + source_wcs: BaseHighLevelWCS | BaseLowLevelWCS, target_wcs: BaseHighLevelWCS | BaseLowLevelWCS, +) -> bool: """ Checks to see if two WCS objects have the same physical types in the same order. @@ -516,10 +535,16 @@ def compare_wcs_physical_types(source_wcs, target_wcs): source_wcs = get_low_level_wcs(source_wcs, 'source_wcs') target_wcs = get_low_level_wcs(target_wcs, 'target_wcs') - return source_wcs.world_axis_physical_types == target_wcs.world_axis_physical_types + return bool(source_wcs.world_axis_physical_types == target_wcs.world_axis_physical_types) -def identify_invariant_axes(source_wcs, target_wcs, input_shape, atol=1e-6, rtol=1e-6): +def identify_invariant_axes( + source_wcs: BaseHighLevelWCS | BaseLowLevelWCS, + target_wcs: BaseHighLevelWCS | BaseLowLevelWCS, + input_shape: tuple[int, ...], + atol: float = 1e-6, + rtol: float = 1e-6, +) -> list[bool]: """ Performs a pixel to pixel transformation to identify if there are any invariant axes between the given source and target WCS objects. diff --git a/ndcube/version.py b/ndcube/version.py index 515c2f0af..cdaf0392c 100644 --- a/ndcube/version.py +++ b/ndcube/version.py @@ -15,3 +15,5 @@ del warnings version = '0.0.0' + +__all__ = ["version"] diff --git a/ndcube/visualization/base.py b/ndcube/visualization/base.py index d981242b0..a6b96452a 100644 --- a/ndcube/visualization/base.py +++ b/ndcube/visualization/base.py @@ -1,4 +1,5 @@ import abc +from typing import Any class BasePlotter(abc.ABC): @@ -6,11 +7,11 @@ class BasePlotter(abc.ABC): Base class for NDCube plotter objects. """ - def __init__(self, ndcube=None): + def __init__(self, ndcube: Any = None) -> None: self._ndcube = ndcube @abc.abstractmethod - def plot(self, *args, **kwargs): + def plot(self, *args: Any, **kwargs: Any) -> Any: """ The default plot method. diff --git a/ndcube/visualization/descriptor.py b/ndcube/visualization/descriptor.py index f0cd5436d..81930cc13 100644 --- a/ndcube/visualization/descriptor.py +++ b/ndcube/visualization/descriptor.py @@ -1,4 +1,5 @@ import functools +from typing import Any, Literal MISSING_MATPLOTLIB_ERROR_MSG = ("matplotlib cannot be imported, so the default plotting " "functionality is disabled. Please install matplotlib") @@ -9,10 +10,10 @@ class PlotterDescriptor: - def __init__(self, default_type=None): + def __init__(self, default_type: 'Literal["mpl_plotter", "mpl_sequence_plotter"] | type[Any] | None' = None) -> None: self._default_type = default_type - def __set_name__(self, owner, name): + def __set_name__(self, owner: type, name: str) -> None: """ This function is called when the class the descriptor is attached to is initialized. @@ -26,9 +27,9 @@ def __set_name__(self, owner, name): self._attribute_name = f"_{name}" plotter = self._resolve_default_type(raise_error=False) if plotter is not None and hasattr(plotter, "plot"): - functools.update_wrapper(owner.plot, plotter.plot) + functools.update_wrapper(owner.plot, plotter.plot) # type: ignore[attr-defined] - def _resolve_default_type(self, raise_error=True): + def _resolve_default_type(self, raise_error: bool = True) -> "type[Any] | None": # We special case the default MatplotlibPlotter so that we can # delay the import of matplotlib until the plotter is first # accessed. @@ -52,12 +53,14 @@ def _resolve_default_type(self, raise_error=True): return None if self._default_type is not None: - return self._default_type + # `in` doesn't narrow Literal types like `==` does, but the block above + # always returns when self._default_type is one of the two literal strings. + return self._default_type # type: ignore[return-value] # pyright: ignore[reportReturnType] # If we have no default type then just return None return None - def __get__(self, obj, objtype=None): + def __get__(self, obj: Any, objtype: type | None = None) -> Any: if obj is None: return None @@ -70,8 +73,9 @@ def __get__(self, obj, objtype=None): return getattr(obj, self._attribute_name) - def __set__(self, obj, value): - if not isinstance(value, type): + def __set__(self, obj: Any, value: "type[Any]") -> None: + # Runtime guard against callers not respecting the type hints. + if not isinstance(value, type): # pyright: ignore[reportUnnecessaryIsInstance] raise TypeError( "Plotter attribute can only be set with an uninitialised plotter object.") diff --git a/ndcube/visualization/mpl_plotter.py b/ndcube/visualization/mpl_plotter.py index d3e6d1345..58e9d0220 100644 --- a/ndcube/visualization/mpl_plotter.py +++ b/ndcube/visualization/mpl_plotter.py @@ -1,4 +1,5 @@ import warnings +from typing import Any import matplotlib.pyplot as plt import numpy as np @@ -20,8 +21,8 @@ class MatplotlibPlotter(BasePlotter): Provide visualization methods for NDCube which use `matplotlib`. """ - def plot(self, axes=None, plot_axes=None, axes_coordinates=None, - axes_units=None, data_unit=None, wcs=None, **kwargs): + def plot(self, axes: Any = None, plot_axes: Any = None, axes_coordinates: Any = None, + axes_units: Any = None, data_unit: Any = None, wcs: Any = None, **kwargs: Any) -> Any: """ Visualize the `~ndcube.NDCube`. @@ -92,14 +93,14 @@ def plot(self, axes=None, plot_axes=None, axes_coordinates=None, return ax - def _not_visible_coords(self, axes, axes_coordinates): + def _not_visible_coords(self, axes: Any, axes_coordinates: Any) -> set[Any]: """ Based on an axes object and axes_coords, work out which coords should not be visible. """ visible_coords = {item[1] for item in axes.coords._aliases.items() if item[0] in axes_coordinates} return set(axes.coords._aliases.values()).difference(visible_coords) - def _apply_axes_coordinates(self, axes, axes_coordinates): + def _apply_axes_coordinates(self, axes: Any, axes_coordinates: Any) -> None: """ Hide ticks and labels for non-visible axes based on axes_coordinates. """ @@ -107,8 +108,8 @@ def _apply_axes_coordinates(self, axes, axes_coordinates): axes.coords[coord_index].set_ticks_visible(False) axes.coords[coord_index].set_ticklabel_visible(False) - def _plot_1D_cube(self, wcs, axes=None, axes_coordinates=None, axes_units=None, - data_unit=None, **kwargs): + def _plot_1D_cube(self, wcs: Any, axes: Any = None, axes_coordinates: Any = None, axes_units: Any = None, + data_unit: Any = None, **kwargs: Any) -> Any: if axes is None: axes = plt.subplot(projection=wcs) @@ -157,8 +158,8 @@ def _plot_1D_cube(self, wcs, axes=None, axes_coordinates=None, axes_units=None, return axes - def _plot_2D_cube(self, wcs, axes=None, plot_axes=None, axes_coordinates=None, - axes_units=None, data_unit=None, **kwargs): + def _plot_2D_cube(self, wcs: Any, axes: Any = None, plot_axes: Any = None, axes_coordinates: Any = None, + axes_units: Any = None, data_unit: Any = None, **kwargs: Any) -> Any: if axes is None: axes = plt.subplot(projection=wcs, slices=plot_axes) @@ -193,8 +194,8 @@ def _plot_2D_cube(self, wcs, axes=None, plot_axes=None, axes_coordinates=None, return axes - def _animate_cube(self, wcs, plot_axes=None, axes_coordinates=None, - axes_units=None, data_unit=None, **kwargs): + def _animate_cube(self, wcs: Any, plot_axes: Any = None, axes_coordinates: Any = None, + axes_units: Any = None, data_unit: Any = None, **kwargs: Any) -> Any: try: from mpl_animators import ArrayAnimatorWCS # noqa: PLC0415 except ImportError as e: @@ -222,7 +223,7 @@ def _animate_cube(self, wcs, plot_axes=None, axes_coordinates=None, return ax - def _as_mpl_axes(self): + def _as_mpl_axes(self) -> tuple[Any, dict[str, Any]]: """ Compatibility hook for Matplotlib and WCSAxes. This functionality requires the WCSAxes package to work. The reason @@ -234,13 +235,14 @@ def _as_mpl_axes(self): and this will generate a plot with the correct WCS coordinates on the axes. See https://wcsaxes.readthedocs.io for more information. """ - kwargs = {'wcs': self._ndcube.wcs} + kwargs: dict[str, Any] = {'wcs': self._ndcube.wcs} n_dim = len(self._ndcube.shape) if n_dim > 2: kwargs['slices'] = ['x', 'y'] + [None] * (n_dim - 2) return WCSAxes, kwargs - def _prep_animate_args(self, wcs, plot_axes, axes_units, data_unit): + def _prep_animate_args(self, wcs: Any, plot_axes: Any, axes_units: Any, + data_unit: Any) -> tuple[Any, Any, list[Any], dict[str, Any]]: # If data_unit set, convert data to that unit if data_unit is None: data = self._ndcube.data @@ -251,7 +253,7 @@ def _prep_animate_args(self, wcs, plot_axes, axes_units, data_unit): if self._ndcube.mask is not None: data = np.ma.masked_array(data, self._ndcube.mask) - coord_params = {} + coord_params: dict[str, Any] = {} if axes_units is not None: for axis_unit, coord_name in zip(axes_units, wcs.world_axis_physical_types): coord_params[coord_name] = {'format_unit': axis_unit} diff --git a/ndcube/visualization/mpl_sequence_plotter.py b/ndcube/visualization/mpl_sequence_plotter.py index 1b1e477c3..eac9d1c7f 100644 --- a/ndcube/visualization/mpl_sequence_plotter.py +++ b/ndcube/visualization/mpl_sequence_plotter.py @@ -1,3 +1,5 @@ +from typing import Any + from mpl_animators import ArrayAnimatorWCS from astropy.wcs.wcsapi import BaseLowLevelWCS @@ -16,7 +18,7 @@ class MatplotlibSequencePlotter(BasePlotter): which is assumed to employ the `~ndcube.visualization.mpl_plotter.MatplotlibPlotter`. """ - def plot(self, sequence_axis_coords=None, sequence_axis_unit=None, **kwargs): + def plot(self, sequence_axis_coords: Any = None, sequence_axis_unit: Any = None, **kwargs: Any) -> Any: """ Visualize the `~ndcube.NDCubeSequence`. @@ -34,7 +36,7 @@ def plot(self, sequence_axis_coords=None, sequence_axis_unit=None, **kwargs): raise NotImplementedError("Visualizing sequences of 1-D cubes not currently supported.") return self.animate(sequence_axis_coords, sequence_axis_unit, **kwargs) - def animate(self, sequence_axis_coords=None, sequence_axis_unit=None, **kwargs): + def animate(self, sequence_axis_coords: Any = None, sequence_axis_unit: Any = None, **kwargs: Any) -> "SequenceAnimator": """ Animate the `~ndcube.NDCubeSequence` with the sequence axis as a slider. @@ -57,7 +59,7 @@ def animate(self, sequence_axis_coords=None, sequence_axis_unit=None, **kwargs): return SequenceAnimator(self._ndcube, sequence_axis_coords=sequence_axis_coords, sequence_axis_unit=sequence_axis_unit, **kwargs) -class SequenceAnimator(ArrayAnimatorWCS): +class SequenceAnimator(ArrayAnimatorWCS): # type: ignore[misc] """ Animate an NDCubeSequence of NDCubes with >1 dimension. @@ -79,7 +81,8 @@ class SequenceAnimator(ArrayAnimatorWCS): The unit in which to display the sequence_axis_coords. """ - def __init__(self, sequence, sequence_axis_coords=None, sequence_axis_unit=None, **kwargs): + def __init__(self, sequence: Any, sequence_axis_coords: Any = None, sequence_axis_unit: Any = None, + **kwargs: Any) -> None: if sequence_axis_coords is not None: raise NotImplementedError("Setting sequence_axis_coords not yet supported.") if sequence_axis_unit is not None: @@ -100,7 +103,7 @@ def __init__(self, sequence, sequence_axis_coords=None, sequence_axis_unit=None, n_cube_dims, init_wcs, plot_axes, axes_coordinates, axes_units) # Define sequence axis slider properties and add to kwargs. - base_kwargs = {"slider_functions": [self._sequence_slider_function], + base_kwargs: dict[str, Any] = {"slider_functions": [self._sequence_slider_function], "slider_ranges": [[0, len(self._cubes)]]} base_kwargs.update(kwargs) @@ -111,7 +114,7 @@ def __init__(self, sequence, sequence_axis_coords=None, sequence_axis_unit=None, wcs = wcs.low_level_wcs super().__init__(data, wcs, plot_axes, coord_params=coord_params, **base_kwargs) - def _sequence_slider_function(self, val, artist, slider): + def _sequence_slider_function(self, val: Any, artist: Any, slider: Any) -> None: self._sequence_idx = int(val) self.data, self.wcs, _, _ = self._cubes[self._sequence_idx].plotter._prep_animate_args( self._cubes[self._sequence_idx].wcs, self._plot_axes, self._axes_units, self._data_unit) diff --git a/ndcube/visualization/plotting_utils.py b/ndcube/visualization/plotting_utils.py index b9d982611..b2d1d18a4 100644 --- a/ndcube/visualization/plotting_utils.py +++ b/ndcube/visualization/plotting_utils.py @@ -1,9 +1,12 @@ +from typing import Any + import astropy.units as u __all__ = ['prep_plot_kwargs', 'set_wcsaxes_format_units'] -def _expand_ellipsis(ndim, plist): +def _expand_ellipsis(ndim: int, plist: Any) -> list[Any]: + plist = list(plist) if Ellipsis in plist: if plist.count(Ellipsis) > 1: raise IndexError("Only single ellipsis ('...') is permitted.") @@ -19,7 +22,8 @@ def _expand_ellipsis(ndim, plist): return plist -def _expand_ellipsis_axis_coordinates(plist, wapt): +def _expand_ellipsis_axis_coordinates(plist: Any, wapt: Any) -> list[Any]: + plist = list(plist) if Ellipsis in plist: if plist.count(Ellipsis) > 1: raise IndexError("Only single ellipsis ('...') is permitted.") @@ -35,7 +39,8 @@ def _expand_ellipsis_axis_coordinates(plist, wapt): return plist -def prep_plot_kwargs(naxis, wcs, plot_axes, axes_coordinates, axes_units): +def prep_plot_kwargs(naxis: int, wcs: Any, plot_axes: Any, axes_coordinates: Any, + axes_units: Any) -> tuple[list[Any], list[Any] | None, list[Any] | None]: """ Prepare the kwargs for the plotting functions. @@ -85,7 +90,7 @@ def prep_plot_kwargs(naxis, wcs, plot_axes, axes_coordinates, axes_units): return plot_axes, axes_coordinates, axes_units -def set_wcsaxes_format_units(coord_map, wcs, axes_units=None): +def set_wcsaxes_format_units(coord_map: Any, wcs: Any, axes_units: list[Any] | None = None) -> None: """ Given an `~astropy.visualization.wcsaxes.coordinates_map.CoordinatesMap` object set the format units. diff --git a/ndcube/wcs/tools.py b/ndcube/wcs/tools.py index 14b80a20b..396f6611b 100644 --- a/ndcube/wcs/tools.py +++ b/ndcube/wcs/tools.py @@ -1,9 +1,11 @@ from numbers import Integral +from collections.abc import Iterable, Sequence import numpy as np +import numpy.typing as npt from astropy.wcs import WCS -from astropy.wcs.wcsapi import SlicedLowLevelWCS +from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS, SlicedLowLevelWCS from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper from ndcube.wcs.wrappers import ResampledLowLevelWCS @@ -11,7 +13,7 @@ __all__ = ["unwrap_wcs_to_fitswcs"] -def unwrap_wcs_to_fitswcs(wcs): +def unwrap_wcs_to_fitswcs(wcs: BaseLowLevelWCS | BaseHighLevelWCS) -> tuple[WCS, npt.NDArray[np.bool_]]: """ Create FITS-WCS equivalent to (nested) WCS wrapper object. @@ -24,8 +26,8 @@ def unwrap_wcs_to_fitswcs(wcs): Parameters ---------- - wcs: `~astropy.wcs.wcsapi.BaseWCSWrapper` - The WCS Wrapper object. + wcs: `~astropy.wcs.wcsapi.BaseLowLevelWCS` or `~astropy.wcs.wcsapi.BaseHighLevelWCS` + The (possibly wrapped) WCS object, low- or high-level. Base level WCS implementation must be FITS-WCS. Returns @@ -62,15 +64,16 @@ def unwrap_wcs_to_fitswcs(wcs): factor = np.ones(fitswcs.naxis) offset = np.zeros(fitswcs.naxis) kept_wcs_axes = dropped_data_axes[::-1] == False # WCS-order # NOQA: E712 - factor[kept_wcs_axes] = low_level_wrapper._factor - offset[kept_wcs_axes] = low_level_wrapper._offset + factor[kept_wcs_axes] = low_level_wrapper._factor # pyright: ignore[reportPrivateUsage] + offset[kept_wcs_axes] = low_level_wrapper._offset # pyright: ignore[reportPrivateUsage] fitswcs = _resample_fitswcs(fitswcs, factor, offset) else: raise TypeError("Unrecognized/unsupported WCS Wrapper type: {type(low_level_wrapper)}") return fitswcs, dropped_data_axes -def _slice_fitswcs(fitswcs, slice_items, numpy_order=True, shape=None): +def _slice_fitswcs(fitswcs: WCS, slice_items: Iterable[slice | int], numpy_order: bool = True, + shape: Sequence[int] | None = None) -> tuple[WCS, npt.NDArray[np.bool_]]: """ Slice a FITS-WCS. @@ -82,7 +85,7 @@ def _slice_fitswcs(fitswcs, slice_items, numpy_order=True, shape=None): ---------- fitswcs: `astropy.wcs.WCS` The FITS-WCS object to be sliced. - slice_items: iterable of `slice` objects or `int` + slice_items: iterable of ``slice`` objects or `int` The slices to by applied to each axis. If an `int` is provided, the axis is sliced to length-1, but not dropped. However, its corresponding entry in the ``dropped_data_axes`` output is marked True. @@ -102,16 +105,16 @@ def _slice_fitswcs(fitswcs, slice_items, numpy_order=True, shape=None): Denotes which axes must have been dropped from the data array by slicing wrappers. Order of axes (numpy or WCS) is dictated by ``numpy_order`` kwarg. """ - def negative_index_error_msg(x): return ( + def negative_index_error_msg(x: int) -> str: return ( f"Negative indexing not supported as {x}th axis length is 0 in " "underlying FITS-WCS. Supply axes lengths via shape kwarg.") naxis = fitswcs.naxis dropped_data_axes = np.zeros(naxis, dtype=bool) # Sanitize inputs if shape is None: - shape = fitswcs._naxis + axis_shape = fitswcs._naxis if numpy_order: - shape = shape[::-1] + axis_shape = axis_shape[::-1] else: if len(shape) != naxis: raise ValueError("shape kwarg must be same length as number of pixel axes " @@ -119,8 +122,9 @@ def negative_index_error_msg(x): return ( if not all(isinstance(s, Integral) for s in shape): raise TypeError("All elements of ``shape`` must be integers. " f"shapes types = {[type(s) for s in shape]}") + axis_shape = shape slice_items = list(slice_items) - for i, (item, len_axis) in enumerate(zip(slice_items, shape)): + for i, (item, len_axis) in enumerate(zip(slice_items, axis_shape)): if isinstance(item, Integral): # Mark axis corresponding to int item as dropped from data array. dropped_data_axes[i] = True @@ -138,8 +142,9 @@ def negative_index_error_msg(x): return ( if start_neg or stop_neg: if len_axis == 0: raise ValueError(negative_index_error_msg(i)) - start = len_axis + item.start if start_neg else item.start - stop = len_axis + item.stop if stop_neg else item.stop + # start_neg/stop_neg guarantee item.start/item.stop aren't None here. + start = len_axis + item.start if start_neg else item.start # pyright: ignore[reportOperatorIssue] + stop = len_axis + item.stop if stop_neg else item.stop # pyright: ignore[reportOperatorIssue] slice_items[i] = slice(start, stop, item.step) else: raise TypeError("All slice_items must be a slice or an int. " @@ -149,7 +154,7 @@ def negative_index_error_msg(x): return ( return sliced_wcs, dropped_data_axes -def _resample_fitswcs(fitswcs, factor, offset=0): +def _resample_fitswcs(fitswcs: WCS, factor: npt.ArrayLike, offset: npt.ArrayLike = 0) -> WCS: """ Resample the plate scale of a FITS-WCS by a given factor. diff --git a/ndcube/wcs/wrappers/compound_wcs.py b/ndcube/wcs/wrappers/compound_wcs.py index 6f0af26b5..af1bd81c8 100644 --- a/ndcube/wcs/wrappers/compound_wcs.py +++ b/ndcube/wcs/wrappers/compound_wcs.py @@ -1,13 +1,17 @@ +from typing import Any from functools import reduce +from collections.abc import Iterable import numpy as np +import numpy.typing as npt +from astropy.wcs.wcsapi import BaseLowLevelWCS from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper __all__ = ['CompoundLowLevelWCS'] -def tuplesum(lists): +def tuplesum(lists: Iterable[Iterable[Any]]) -> tuple[Any, ...]: return reduce(tuple.__add__, map(tuple, lists)) @@ -27,25 +31,25 @@ class Mapping: """ - def __init__(self, mapping): + def __init__(self, mapping: tuple[int, ...]) -> None: self.mapping = mapping self.n_inputs = max(mapping) + 1 self.n_outputs = len(mapping) - def __call__(self, *values): + def __call__(self, *values: Any) -> tuple[Any, ...]: return tuple(values[idx] for idx in self.mapping) @property - def inverse(self): + def inverse(self) -> "Mapping": mapping = tuple(self.mapping.index(idx) for idx in range(self.n_inputs)) return type(self)(mapping) - def __repr__(self): + def __repr__(self) -> str: return f'' -class CompoundLowLevelWCS(BaseWCSWrapper): +class CompoundLowLevelWCS(BaseWCSWrapper): # type: ignore[misc] """ A wrapper that takes multiple low level WCS objects and makes a compound WCS that combines them. @@ -69,7 +73,8 @@ class CompoundLowLevelWCS(BaseWCSWrapper): ``world_to_pixel`` are the same from all WCSes. """ - def __init__(self, *wcs, mapping=None, pixel_atol=1e-8): + def __init__(self, *wcs: BaseLowLevelWCS, mapping: tuple[int, ...] | None = None, + pixel_atol: float = 1e-8) -> None: self._wcs = wcs if not mapping: @@ -87,26 +92,26 @@ def __init__(self, *wcs, mapping=None, pixel_atol=1e-8): self.pixel_shape @property - def _all_pixel_n_dim(self): + def _all_pixel_n_dim(self) -> int: return sum([w.pixel_n_dim for w in self._wcs]) @property - def pixel_n_dim(self): + def pixel_n_dim(self) -> int: return self.mapping.n_inputs @property - def world_n_dim(self): + def world_n_dim(self) -> int: return sum([w.world_n_dim for w in self._wcs]) @property - def world_axis_physical_types(self): + def world_axis_physical_types(self) -> tuple[Any, ...]: return tuplesum([w.world_axis_physical_types for w in self._wcs]) @property - def world_axis_units(self): + def world_axis_units(self) -> tuple[Any, ...]: return tuplesum([w.world_axis_units for w in self._wcs]) - def pixel_to_world_values(self, *pixel_arrays): + def pixel_to_world_values(self, *pixel_arrays: Any) -> tuple[Any, ...]: pixel_arrays = self.mapping(*pixel_arrays) world_arrays = [] for w in self._wcs: @@ -119,7 +124,7 @@ def pixel_to_world_values(self, *pixel_arrays): pixel_arrays = pixel_arrays[w.pixel_n_dim:] return tuple(world_arrays) - def world_to_pixel_values(self, *world_arrays): + def world_to_pixel_values(self, *world_arrays: Any) -> tuple[Any, ...]: pixel_arrays = [] for w in self._wcs: world_arrays_sub = world_arrays[:w.world_n_dim] @@ -130,9 +135,15 @@ def world_to_pixel_values(self, *world_arrays): else: pixel_arrays.append(pixel_arrays_sub) + # NOTE: self.mapping.mapping is a plain tuple, so `== mapped_axis` here compares + # tuple-to-int (always False) rather than doing an elementwise numpy comparison. + # This means the shared-pixel-axis consistency check below never actually runs. + # Flagged rather than silently fixed: making it elementwise (e.g. via + # np.array(self.mapping.mapping)) causes test_shared_pixel_axis_compound_3d to fail + # on the ordinary roundtrip case, which needs a deliberate look, not a typing-pass fix. mapped_axes = set(self.mapping.mapping) for mapped_axis in mapped_axes: - idx, = np.atleast_1d(self.mapping.mapping == mapped_axis).nonzero() + idx, = np.atleast_1d(self.mapping.mapping == mapped_axis).nonzero() # type: ignore[comparison-overlap] if len(idx) > 1: idx_0 = idx[0] for idx_n in idx[1:]: @@ -145,7 +156,7 @@ def world_to_pixel_values(self, *world_arrays): return self.mapping.inverse(*pixel_arrays) @property - def world_axis_object_components(self): + def world_axis_object_components(self) -> list[tuple[Any, ...]]: all_components = [] for iw, w in enumerate(self._wcs): all_components += [(f'{component[0]}_{iw}', *component[1:]) for component @@ -153,7 +164,7 @@ def world_axis_object_components(self): return all_components @property - def world_axis_object_classes(self): + def world_axis_object_classes(self) -> dict[str, Any]: # TODO: deal with name conflicts all_classes = {} for iw, w in enumerate(self._wcs): @@ -162,7 +173,7 @@ def world_axis_object_classes(self): return all_classes @property - def pixel_shape(self): + def pixel_shape(self) -> tuple[int, ...] | None: if not any(w.array_shape is None for w in self._wcs): pixel_shape = tuplesum(w.pixel_shape for w in self._wcs) out_shape = self.mapping.inverse(*pixel_shape) @@ -175,7 +186,7 @@ def pixel_shape(self): return None @property - def pixel_bounds(self): + def pixel_bounds(self) -> tuple[tuple[float, float], ...] | None: if any(w.pixel_bounds is not None for w in self._wcs): pixel_bounds = tuplesum(w.pixel_bounds or [() for _ in range(w.pixel_n_dim)] for w in self._wcs) out_bounds = self.mapping.inverse(*pixel_bounds) @@ -190,22 +201,22 @@ def pixel_bounds(self): @property - def pixel_axis_names(self): + def pixel_axis_names(self) -> tuple[str, ...]: pixel_names = tuplesum(w.pixel_axis_names for w in self._wcs) - out_names = self.mapping.inverse(*pixel_names) + out_names = list(self.mapping.inverse(*pixel_names)) for i, ix in enumerate(self.mapping.mapping): if out_names[ix] != pixel_names[i]: out_names[ix] = ' / '.join([out_names[ix], pixel_names[i]]) - return out_names + return tuple(out_names) @property - def world_axis_names(self): + def world_axis_names(self) -> tuple[str, ...]: return tuplesum(w.world_axis_names for w in self._wcs) @property - def axis_correlation_matrix(self): + def axis_correlation_matrix(self) -> npt.NDArray[np.bool_]: full_matrix = np.zeros((self.world_n_dim, self._all_pixel_n_dim), dtype=bool) iw = ip = 0 for w in self._wcs: @@ -220,5 +231,5 @@ def axis_correlation_matrix(self): return matrix @property - def serialized_classes(self): + def serialized_classes(self) -> bool: return any(w.serialized_classes for w in self._wcs) diff --git a/ndcube/wcs/wrappers/reordered_wcs.py b/ndcube/wcs/wrappers/reordered_wcs.py index 84e5565dd..4c00058ef 100644 --- a/ndcube/wcs/wrappers/reordered_wcs.py +++ b/ndcube/wcs/wrappers/reordered_wcs.py @@ -1,13 +1,18 @@ +from typing import Any +from collections.abc import Sequence + import numpy as np +import numpy.typing as npt +from astropy.wcs.wcsapi import BaseLowLevelWCS from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper __all__ = ['ReorderedLowLevelWCS'] -class ReorderedLowLevelWCS(BaseWCSWrapper): +class ReorderedLowLevelWCS(BaseWCSWrapper): # type: ignore[misc] """ A wrapper for a low-level WCS object that has re-ordered pixel and/or world axes. @@ -24,7 +29,7 @@ class ReorderedLowLevelWCS(BaseWCSWrapper): new WCS. """ - def __init__(self, wcs, pixel_order, world_order): + def __init__(self, wcs: BaseLowLevelWCS, pixel_order: Sequence[int], world_order: Sequence[int]) -> None: if sorted(pixel_order) != list(range(wcs.pixel_n_dim)): raise ValueError(f'pixel_order should be a permutation of {list(range(wcs.pixel_n_dim))}') if sorted(world_order) != list(range(wcs.world_n_dim)): @@ -36,47 +41,47 @@ def __init__(self, wcs, pixel_order, world_order): self._world_order_inv = np.argsort(world_order) @property - def world_axis_physical_types(self): + def world_axis_physical_types(self) -> list[Any]: return [self._wcs.world_axis_physical_types[idx] for idx in self._world_order] @property - def world_axis_units(self): + def world_axis_units(self) -> list[Any]: return [self._wcs.world_axis_units[idx] for idx in self._world_order] @property - def pixel_axis_names(self): + def pixel_axis_names(self) -> list[Any]: return [self._wcs.pixel_axis_names[idx] for idx in self._pixel_order] @property - def world_axis_names(self): + def world_axis_names(self) -> list[Any]: return [self._wcs.world_axis_names[idx] for idx in self._world_order] - def pixel_to_world_values(self, *pixel_arrays): - pixel_arrays = [pixel_arrays[idx] for idx in self._pixel_order_inv] - world_arrays = self._wcs.pixel_to_world_values(*pixel_arrays) + def pixel_to_world_values(self, *pixel_arrays: Any) -> list[Any]: + ordered_pixel_arrays = [pixel_arrays[idx] for idx in self._pixel_order_inv] + world_arrays = self._wcs.pixel_to_world_values(*ordered_pixel_arrays) return [world_arrays[idx] for idx in self._world_order] - def world_to_pixel_values(self, *world_arrays): - world_arrays = [world_arrays[idx] for idx in self._world_order_inv] - pixel_arrays = self._wcs.world_to_pixel_values(*world_arrays) + def world_to_pixel_values(self, *world_arrays: Any) -> list[Any]: + ordered_world_arrays = [world_arrays[idx] for idx in self._world_order_inv] + pixel_arrays = self._wcs.world_to_pixel_values(*ordered_world_arrays) return [pixel_arrays[idx] for idx in self._pixel_order] @property - def world_axis_object_components(self): + def world_axis_object_components(self) -> list[Any]: return [self._wcs.world_axis_object_components[idx] for idx in self._world_order] @property - def pixel_shape(self): + def pixel_shape(self) -> tuple[Any, ...] | None: if self._wcs.pixel_shape: return tuple([self._wcs.pixel_shape[idx] for idx in self._pixel_order]) return None @property - def pixel_bounds(self): + def pixel_bounds(self) -> tuple[Any, ...] | None: if self._wcs.pixel_bounds: return tuple([self._wcs.pixel_bounds[idx] for idx in self._pixel_order]) return None @property - def axis_correlation_matrix(self): - return self._wcs.axis_correlation_matrix[self._world_order][:, self._pixel_order] + def axis_correlation_matrix(self) -> npt.NDArray[np.bool_]: + return np.asarray(self._wcs.axis_correlation_matrix[self._world_order][:, self._pixel_order]) diff --git a/ndcube/wcs/wrappers/resampled_wcs.py b/ndcube/wcs/wrappers/resampled_wcs.py index b24805e26..bcc629c41 100644 --- a/ndcube/wcs/wrappers/resampled_wcs.py +++ b/ndcube/wcs/wrappers/resampled_wcs.py @@ -1,12 +1,17 @@ +from typing import Any +from collections.abc import Iterable + import numpy as np +import numpy.typing as npt +from astropy.wcs.wcsapi import BaseLowLevelWCS from astropy.wcs.wcsapi.wrappers.base import BaseWCSWrapper __all__ = ['ResampledLowLevelWCS'] -class ResampledLowLevelWCS(BaseWCSWrapper): +class ResampledLowLevelWCS(BaseWCSWrapper): # type: ignore[misc] """ A wrapper for a low-level WCS object that has down- or up-sampled pixel axes. @@ -76,7 +81,8 @@ class ResampledLowLevelWCS(BaseWCSWrapper): -1 -0.75 -0.5 -0.25 0 0.25 0.5 0.75 1 resampled pixel indices: factor=2, offset=1 """ - def __init__(self, wcs, factor, offset=0): + def __init__(self, wcs: BaseLowLevelWCS, factor: float | Iterable[float], + offset: float | Iterable[float] = 0) -> None: self._wcs = wcs if np.isscalar(factor): factor = [factor] * self.pixel_n_dim @@ -89,36 +95,36 @@ def __init__(self, wcs, factor, offset=0): if len(self._offset) != self.pixel_n_dim: raise ValueError(f"Length of offset must equal number of dimensions {self.pixel_n_dim}.") - def _top_to_underlying_pixels(self, top_pixels): + def _top_to_underlying_pixels(self, top_pixels: npt.NDArray[Any]) -> npt.NDArray[Any]: # Convert user-facing pixel indices to the pixel grid of underlying WCS. factor = self._pad_dims(self._factor, top_pixels.ndim) offset = self._pad_dims(self._offset, top_pixels.ndim) - return (top_pixels + 0.5) * factor - 0.5 + offset + return np.asarray((top_pixels + 0.5) * factor - 0.5 + offset) - def _underlying_to_top_pixels(self, underlying_pixels): + def _underlying_to_top_pixels(self, underlying_pixels: npt.NDArray[Any]) -> npt.NDArray[Any]: # Convert pixel indices of underlying pixel grid to user-facing grid. factor = self._pad_dims(self._factor, underlying_pixels.ndim) offset = self._pad_dims(self._offset, underlying_pixels.ndim) - return (underlying_pixels + 0.5 - offset) / factor - 0.5 + return np.asarray((underlying_pixels + 0.5 - offset) / factor - 0.5) - def _pad_dims(self, arr, ndim): + def _pad_dims(self, arr: npt.NDArray[Any], ndim: int) -> npt.NDArray[Any]: # Pad array with trailing degenerate dimensions. # This make scaling with pixel arrays easier. shape = np.ones(ndim, dtype=int) shape[0] = len(arr) return arr.reshape(tuple(shape)) - def pixel_to_world_values(self, *pixel_arrays): + def pixel_to_world_values(self, *pixel_arrays: Any) -> tuple[Any, ...]: underlying_pixel_arrays = self._top_to_underlying_pixels(np.asarray(pixel_arrays)) - return self._wcs.pixel_to_world_values(*underlying_pixel_arrays) + return self._wcs.pixel_to_world_values(*underlying_pixel_arrays) # type: ignore[no-any-return] - def world_to_pixel_values(self, *world_arrays): + def world_to_pixel_values(self, *world_arrays: Any) -> tuple[Any, ...]: underlying_pixel_arrays = self._wcs.world_to_pixel_values(*world_arrays) top_pixel_arrays = self._underlying_to_top_pixels(np.asarray(underlying_pixel_arrays)) return tuple(array for array in top_pixel_arrays) @property - def pixel_shape(self): + def pixel_shape(self) -> tuple[int | float, ...] | None: # Return pixel shape of resampled grid. # Where shape is an integer, return an int type as its required for some uses. if self._wcs.pixel_shape is None: @@ -131,7 +137,7 @@ def pixel_shape(self): for i, is_int in zip(pixel_shape, int_elements)) @property - def pixel_bounds(self): + def pixel_bounds(self) -> list[tuple[Any, ...]] | None: if self._wcs.pixel_bounds is None: return self._wcs.pixel_bounds top_level_bounds = self._underlying_to_top_pixels(np.asarray(self._wcs.pixel_bounds)) diff --git a/ndcube/wcs/wrappers/tests/test_compound_wcs.py b/ndcube/wcs/wrappers/tests/test_compound_wcs.py index f1b2ba649..a12f527cc 100644 --- a/ndcube/wcs/wrappers/tests/test_compound_wcs.py +++ b/ndcube/wcs/wrappers/tests/test_compound_wcs.py @@ -184,3 +184,33 @@ def test_shared_pixel_axis_compound_3d(spectral_cube_3d_fitswcs, time_1d_fitswcs with pytest.raises(ValueError): wcs.world_to_pixel_values((14, -10, -2.6e+10, -7.0)) + + +class _NamedPixelAxisWCS: + """ + A minimal low-level WCS stub with a non-empty ``pixel_axis_names``. + + FITS WCS objects always report empty-string pixel axis names, so they + can't be used to exercise the case where two WCSes sharing a pixel axis + (via ``mapping``) disagree on that axis's name. + """ + + def __init__(self, pixel_n_dim, pixel_axis_names): + self.pixel_n_dim = pixel_n_dim + self.array_shape = None + self.pixel_bounds = None + self.pixel_axis_names = pixel_axis_names + + +def test_pixel_axis_names_shared_axis_different_names(): + # Regression test: CompoundLowLevelWCS.pixel_axis_names used to build + # out_names via `list(self.mapping.inverse(*pixel_names))` and then try to + # mutate the *tuple* returned by `Mapping.__call__` in place, raising + # "'tuple' object does not support item assignment" whenever a shared + # pixel axis had different names in the contributing WCSes. + wcs_a = _NamedPixelAxisWCS(1, ('a',)) + wcs_b = _NamedPixelAxisWCS(1, ('b',)) + + compound = CompoundLowLevelWCS(wcs_a, wcs_b, mapping=(0, 0)) + + assert compound.pixel_axis_names == ('a / b',) diff --git a/pyproject.toml b/pyproject.toml index 1025186ab..de69d80a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,6 +98,42 @@ exclude = ["ndcube._dev*"] [tool.setuptools_scm] version_file = "ndcube/_version.py" +[tool.mypy] +python_version = "3.12" +strict = true +ignore_missing_imports = true +disable_error_code = "import-untyped" +allow_redefinition = true +exclude = "(^|/)tests/" + +[[tool.mypy.overrides]] +module = ["ndcube.conftest", "*.tests.*"] +ignore_errors = true + +[tool.pyright] +pythonVersion = "3.12" +typeCheckingMode = "strict" + +# Don't infer types from unstubbed library source (astropy/scipy/sunpy have no py.typed); +# treat them as untyped like mypy's ignore_missing_imports does, rather than statically +# analysing their implementation. +useLibraryCodeForTypes = false + +## Stop errors when third-party libraries lack types/stubs +reportMissingImports = "none" +reportMissingTypeStubs = "none" + +# Corrected rules to silence "partially unknown" issues from external packages +reportUnknownVariableType = "none" +reportUnknownMemberType = "none" +reportUnknownArgumentType = "none" +reportUnknownParameterType = "none" + +ignore = [ + "**/tests/**", + "**/conftest.py", +] + [tool.gilesbot] [tool.gilesbot.pull_requests] enabled = true diff --git a/pytest.ini b/pytest.ini index a1c2f7bff..98b346a65 100644 --- a/pytest.ini +++ b/pytest.ini @@ -60,3 +60,7 @@ filterwarnings = ignore:FigureCanvasAgg is non-interactive, and thus cannot be shown:UserWarning # wcsaxes/formatter_locator.py hates angles ignore:.*invalid value encountered in do_format.*:RuntimeWarning + # asdf warning + ignore:.*:asdf.exceptions.AsdfWarning + # astropy + ignore:.*:astropy.utils.exceptions.AstropyPendingDeprecationWarning diff --git a/tox.ini b/tox.ini index e41473dc2..38516f625 100644 --- a/tox.ini +++ b/tox.ini @@ -127,3 +127,25 @@ deps = commands = pre-commit install-hooks pre-commit run --color always --all-files --show-diff-on-failure + +[testenv:mypy] +description = Run mypy type checks +skip_install = false +extras = + plotting + docs +deps = + mypy +commands = + mypy {toxinidir}/ndcube {posargs} + +[testenv:pyright] +description = Run pyright type checks +skip_install = false +extras = + plotting + docs +deps = + pyright +commands = + pyright {toxinidir}/ndcube {posargs} \ No newline at end of file