diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 88e826b7e..de09964c7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,6 +17,7 @@ repos: exclude: dependencies.yaml - id: no-commit-to-branch name: "Don't commit to 'main' directly" + - id: debug-statements - repo: https://github.com/charliermarsh/ruff-pre-commit rev: v0.11.12 diff --git a/src/dxtbx/boost_python/imageset_ext.cc b/src/dxtbx/boost_python/imageset_ext.cc index 858f7b372..07e415efa 100644 --- a/src/dxtbx/boost_python/imageset_ext.cc +++ b/src/dxtbx/boost_python/imageset_ext.cc @@ -40,24 +40,30 @@ namespace dxtbx { namespace boost_python { } } // namespace detail - ImageSetData::masker_ptr make_masker_pointer(boost::python::object masker) { - if (masker == boost::python::object()) { - return ImageSetData::masker_ptr(); - } - return boost::python::extract(masker)(); - } - /** * A constructor for the imageset data class */ std::shared_ptr make_imageset_data1(boost::python::object reader, boost::python::object masker) { - // Create the pointer - std::shared_ptr self( - new ImageSetData(reader, make_masker_pointer(masker))); - - // Return the imageset data - return self; + // Distinguish between the three masker states: + // - A direct GoniometerShadowMasker instance + // - A callable () -> GoniometerShadowMasker (we want to defer render) + // - None + boost::python::extract get_masker_ptr(masker); + ImageSetData *self = nullptr; + if (get_masker_ptr.check()) { + self = new ImageSetData(reader, get_masker_ptr()); + } else if (masker == boost::python::object()) { + self = new ImageSetData(reader, nullptr); + } else if (PyCallable_Check(masker.ptr())) { + self = new ImageSetData(reader, masker); + } else { + PyErr_SetString(PyExc_TypeError, + "Masker object must be: GoniometerShadowMasker | Callable[[], " + "GoniometerShadow] | None"); + throw boost::python::error_already_set(); + } + return std::shared_ptr(self); } /** @@ -70,8 +76,7 @@ namespace dxtbx { namespace boost_python { boost::python::dict params, boost::python::object format) { // Create the pointer - std::shared_ptr self( - new ImageSetData(reader, make_masker_pointer(masker))); + auto self = make_imageset_data1(reader, masker); // Set some stuff self->set_template(filename_template); @@ -473,7 +478,8 @@ namespace dxtbx { namespace boost_python { &ExternalLookupItem::get_filename, &ExternalLookupItem::set_filename) .add_property( - "data", &ExternalLookupItem::get_data, &ExternalLookupItem::set_data); + "data", &ExternalLookupItem::get_data, &ExternalLookupItem::set_data) + .def("set_data_generator", &ExternalLookupItem::set_data_generator); } /** diff --git a/src/dxtbx/dxtbx_format_image_ext.pyi b/src/dxtbx/dxtbx_format_image_ext.pyi index 32d452b93..e7aa92fa4 100644 --- a/src/dxtbx/dxtbx_format_image_ext.pyi +++ b/src/dxtbx/dxtbx_format_image_ext.pyi @@ -1,61 +1,59 @@ from __future__ import annotations -from typing import Any +from collections.abc import Sequence +from typing import Any, Generic, TypeVar, overload import pycbf -class ImageBool: - def __init__(self, *args, **kwargs) -> None: ... +import scitbx.array_family.flex as flex + +T = TypeVar("T") + +class ImageTypedBase(Generic[T]): + """Not a real class, only used here for abstract bases""" + + def __init__(self, data: Sequence[T], /) -> None: ... + def __len__(self) -> int: ... + def empty(self) -> bool: ... + def n_tiles(self) -> int: ... + +class ImageBool(ImageTypedBase[bool]): + @overload + def __init__(self, data: flex.bool, /) -> None: ... def append(self, *args, **kwargs) -> Any: ... - def empty(dxtbx) -> Any: ... - def n_tiles(dxtbx) -> Any: ... def tile(dxtbx, unsignedlong) -> Any: ... def tile_names(dxtbx) -> Any: ... def __getitem__(dxtbx, unsignedlong) -> Any: ... - def __getstate__(dxtbx) -> Any: ... def __iter__(boost) -> Any: ... - def __len__(dxtbx) -> Any: ... - def __reduce__(self) -> Any: ... - def __setstate__(dxtbx, boost) -> Any: ... -class ImageBuffer: - def __init__(self, *args, **kwargs) -> None: ... - def as_double(dxtbx) -> Any: ... - def as_float(dxtbx) -> Any: ... - def as_int(dxtbx) -> Any: ... - def is_double(dxtbx) -> Any: ... - def is_empty(dxtbx) -> Any: ... - def is_float(dxtbx) -> Any: ... - def is_int(dxtbx) -> Any: ... - def __reduce__(self) -> Any: ... - -class ImageDouble: - def __init__(self, *args, **kwargs) -> None: ... +class ImageDouble(ImageTypedBase[float]): + @overload + def __init__(self, data: flex.double, /) -> None: ... def append(self, *args, **kwargs) -> Any: ... - def empty(dxtbx) -> Any: ... - def n_tiles(dxtbx) -> Any: ... def tile(dxtbx, unsignedlong) -> Any: ... def tile_names(dxtbx) -> Any: ... def __getitem__(dxtbx, unsignedlong) -> Any: ... - def __getstate__(dxtbx) -> Any: ... def __iter__(boost) -> Any: ... - def __len__(dxtbx) -> Any: ... - def __reduce__(self) -> Any: ... - def __setstate__(dxtbx, boost) -> Any: ... -class ImageInt: - def __init__(self, *args, **kwargs) -> None: ... +class ImageInt(ImageTypedBase[int]): + @overload + def __init__(self, data: flex.int, /) -> None: ... def append(self, *args, **kwargs) -> Any: ... - def empty(dxtbx) -> Any: ... - def n_tiles(dxtbx) -> Any: ... def tile(dxtbx, unsignedlong) -> Any: ... def tile_names(dxtbx) -> Any: ... def __getitem__(dxtbx, unsignedlong) -> Any: ... - def __getstate__(dxtbx) -> Any: ... def __iter__(boost) -> Any: ... - def __len__(dxtbx) -> Any: ... + +class ImageBuffer: + def __init__(self, *args, **kwargs) -> None: ... + def as_double(dxtbx) -> Any: ... + def as_float(dxtbx) -> Any: ... + def as_int(dxtbx) -> Any: ... + def is_double(dxtbx) -> Any: ... + def is_empty(dxtbx) -> Any: ... + def is_float(dxtbx) -> Any: ... + def is_int(dxtbx) -> Any: ... def __reduce__(self) -> Any: ... - def __setstate__(dxtbx, boost) -> Any: ... class ImageTileBool: def __init__(self, *args, **kwargs) -> None: ... diff --git a/src/dxtbx/dxtbx_imageset_ext.pyi b/src/dxtbx/dxtbx_imageset_ext.pyi index 9debeae44..445aeac30 100644 --- a/src/dxtbx/dxtbx_imageset_ext.pyi +++ b/src/dxtbx/dxtbx_imageset_ext.pyi @@ -1,36 +1,47 @@ from __future__ import annotations -from typing import Any, Union +from typing import Any, Callable, Protocol, Type, Union, overload from scitbx.array_family import flex +from dxtbx.format.Format import Format +from dxtbx.format.image import ImageBool, ImageDouble +from dxtbx.masking import GoniometerShadowMasker +from dxtbx.model import Beam, Detector, Goniometer, Scan + +class Reader(Protocol): + def read(self, index: int | None) -> Any: ... + def is_single_file_reader(self) -> bool: ... + def paths(self) -> list[str]: ... + def master_path(self) -> str: ... + def identifiers(self) -> list[str]: ... + def __len__(self) -> int: ... + ImageData = Union[flex.int, flex.double, flex.float] class ExternalLookup: - def __init__(self, *args, **kwargs) -> None: ... - def __reduce__(self) -> Any: ... @property - def dx(self) -> Any: ... + def dx(self) -> ExternalLookupItemDouble: ... @property - def dy(self) -> Any: ... + def dy(self) -> ExternalLookupItemDouble: ... @property - def gain(self) -> Any: ... + def gain(self) -> ExternalLookupItemDouble: ... @property - def mask(self) -> Any: ... + def mask(self) -> ExternalLookupItemBool: ... @property - def pedestal(self) -> Any: ... + def pedestal(self) -> ExternalLookupItemDouble: ... class ExternalLookupItemBool: - data: Any - filename: Any - def __init__(self, *args, **kwargs) -> None: ... - def __reduce__(self) -> Any: ... + data: ImageBool + filename: str + def set_data_generator(self, generator: Callable[[], ImageBool | None]) -> None: ... class ExternalLookupItemDouble: - data: Any - filename: Any - def __init__(self, *args, **kwargs) -> None: ... - def __reduce__(self) -> Any: ... + data: ImageDouble + filename: str + def set_data_generator( + self, generator: Callable[[], ImageDouble | None] + ) -> None: ... class ImageGrid(ImageSet): def __init__(self, *args, **kwargs) -> None: ... @@ -40,7 +51,25 @@ class ImageGrid(ImageSet): def __reduce__(self) -> Any: ... class ImageSequence(ImageSet): - def __init__(self, *args, **kwargs) -> None: ... + @overload + def __init__( + self, + data: ImageSetData, + beam: Beam, + detector: Detector, + goniometer: Goniometer, + scan: Scan, + ) -> None: ... + @overload + def __init__( + self, + data: ImageSetData, + beam: Beam, + detector: Detector, + goniometer: Goniometer, + scan: Scan, + indices: flex.size_t, + ) -> None: ... def complete_set(self) -> Any: ... def get_array_range(self) -> Any: ... def get_beam(self) -> Any: ... @@ -90,38 +119,42 @@ class ImageSet: def __ne__(self, other) -> Any: ... def __reduce__(self) -> Any: ... @property - def external_lookup(self) -> Any: ... + def external_lookup(self) -> ExternalLookup: ... class ImageSetData: - def __init__(self, *args, **kwargs) -> None: ... - def get_beam(self, int) -> Any: ... - def get_data(self, int) -> Any: ... - def get_detector(self, int) -> Any: ... - def get_format_class(self) -> Any: ... - def get_goniometer(self, int) -> Any: ... - def get_image_identifier(self, int) -> Any: ... - def get_master_path(self) -> Any: ... - def get_params(self) -> Any: ... - def get_path(self, int) -> Any: ... - def get_scan(self, int) -> Any: ... - def get_template(self) -> Any: ... - def get_vendor(self) -> Any: ... - def has_single_file_reader(self) -> Any: ... - def is_marked_for_rejection(self, int) -> Any: ... - def mark_for_rejection(self, int, bool) -> Any: ... - def masker(self) -> Any: ... - def reader(self) -> Any: ... - def set_beam(self, boost, int) -> Any: ... - def set_detector(self, boost, int) -> Any: ... - def set_format_class(self, boost) -> Any: ... - def set_goniometer(self, boost, int) -> Any: ... - def set_params(self, boost) -> Any: ... - def set_scan(self, boost, int) -> Any: ... - def set_template(self, *args, **kwargs) -> Any: ... - def set_vendor(self, *args, **kwargs) -> Any: ... - def __getinitargs__(self) -> Any: ... - def __getstate__(self) -> Any: ... - def __reduce__(self) -> Any: ... - def __setstate__(self, boost) -> Any: ... + def __init__( + self, + reader: Reader, + masker: GoniometerShadowMasker | Callable[[], GoniometerShadowMasker] | None, + template: str = "", + vendor: str = "", + params: dict | None = None, + format: Format | None = None, + ) -> None: ... + def get_beam(self, index: int) -> int: ... + def get_data(self, index: int) -> int: ... + def get_detector(self, index: int) -> int: ... + def get_goniometer(self, index: int) -> int: ... + def get_scan(self, index: int) -> int: ... + def get_format_class(self) -> Type[Format]: ... + def get_image_identifier(self, index: int) -> str: ... + def get_params(self) -> dict | None: ... + def get_path(self, int) -> str: ... + def get_master_path(self) -> str: ... + def get_template(self) -> str: ... + def get_vendor(self) -> str: ... + def has_single_file_reader(self) -> bool: ... + def is_marked_for_rejection(self, index: int) -> bool: ... + def mark_for_rejection(self, index: bool, is_rejected: bool) -> None: ... + def masker(self) -> GoniometerShadowMasker | None: ... + def reader(self) -> Reader: ... + def set_beam(self, value: int, index: int) -> None: ... + def set_detector(self, value: int, index: int) -> None: ... + def set_goniometer(self, value: int, index: int) -> None: ... + def set_scan(self, value: int, index: int) -> None: ... + def set_params(self, value: dict | None) -> None: ... + def set_format_class(self, format: Type[Format]) -> None: ... + def set_template(self, template: str) -> None: ... + def set_vendor(self, vendor: str) -> None: ... @property - def external_lookup(self) -> Any: ... + def external_lookup(self) -> ExternalLookup: ... diff --git a/src/dxtbx/dxtbx_model_ext.pyi b/src/dxtbx/dxtbx_model_ext.pyi index c9fe92c74..0d9766333 100644 --- a/src/dxtbx/dxtbx_model_ext.pyi +++ b/src/dxtbx/dxtbx_model_ext.pyi @@ -389,6 +389,10 @@ class ExperimentList: ) -> flex.size_t: ... def is_consistent(self) -> bool: ... def __len__(self) -> int: ... + # Injected by python at runtime + def as_file(self, filename: str, **kwargs): + """Dump experiment list as file.""" + ... class GoniometerBase: pass diff --git a/src/dxtbx/format/Format.py b/src/dxtbx/format/Format.py index 397e8cfa3..6731cdfad 100644 --- a/src/dxtbx/format/Format.py +++ b/src/dxtbx/format/Format.py @@ -12,14 +12,17 @@ import bz2 import functools import os -from collections.abc import Callable +from collections.abc import Callable, Sequence from io import IOBase -from typing import ClassVar +from typing import TYPE_CHECKING, ClassVar, Literal, overload import libtbx +import scitbx.array_family.flex as flex import dxtbx.filecache_controller +from dxtbx import model from dxtbx.format.image import ImageBool +from dxtbx.imageset import ImageSequence, ImageSet, ImageSetLazy from dxtbx.model import MultiAxisGoniometer from dxtbx.model.beam import BeamFactory from dxtbx.model.detector import DetectorFactory @@ -33,6 +36,8 @@ except ImportError: gzip = None # type: ignore +if TYPE_CHECKING: + from dxtbx.masking import GoniometerShadowMasker _cache_controller = dxtbx.filecache_controller.simple_controller() @@ -284,7 +289,9 @@ def get_reader(cls): """ return functools.partial(Reader, cls) - def get_masker(self, goniometer=None): + def get_masker( + self, goniometer: model.Goniometer | None = None + ) -> GoniometerShadowMasker | None: """ Return a masker class """ @@ -298,32 +305,67 @@ def get_masker(self, goniometer=None): masker = None return masker + @overload @classmethod def get_imageset( - Class, - input_filenames, - beam=None, - detector=None, - goniometer=None, - scan=None, - as_imageset=False, - as_sequence=False, - single_file_indices=None, - format_kwargs=None, - template=None, - check_format=True, - ): + cls, + filenames: str | Sequence[str], + beam: model.Beam | None = None, + detector: model.Detector | None = None, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + as_imageset: Literal[False] = ..., + as_sequence: Literal[True] = ..., + single_file_indices: Sequence[int] | None = None, + format_kwargs: dict | None = None, + template: str | None = None, + check_format: bool = True, + ) -> ImageSequence: + pass + + @overload + @classmethod + def get_imageset( + cls, + filenames: str | Sequence[str], + beam: model.Beam | None = None, + detector: model.Detector | None = None, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + as_imageset: bool = False, + as_sequence: bool = False, + single_file_indices: Sequence[int] | None = None, + format_kwargs: dict | None = None, + template: str | None = None, + check_format: bool = True, + ) -> ImageSet | ImageSequence | ImageSetLazy: ... + + @classmethod + def get_imageset( + cls, + filenames: str | Sequence[str], + beam: model.Beam | None = None, + detector: model.Detector | None = None, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + as_imageset: bool = False, + as_sequence: bool = False, + single_file_indices: Sequence[int] | None = None, + format_kwargs: dict | None = None, + template: str | None = None, + check_format: bool = True, + ) -> ImageSet | ImageSequence | ImageSetLazy: """ Factory method to create an imageset """ # Import here to avoid cyclic imports - from dxtbx.imageset import ImageSequence, ImageSet, ImageSetData + from dxtbx.imageset import ImageSetData # Turn entries that are filenames into absolute paths filenames = [ os.fspath(os.path.abspath(x)) if not get_url_scheme(x) else x - for x in input_filenames + for x in filenames ] # Make it a dict @@ -331,13 +373,13 @@ def get_imageset( format_kwargs = {} # Get some information from the format class - reader = Class.get_reader()(filenames, **format_kwargs) + reader = cls.get_reader()(filenames, **format_kwargs) # Get the format instance if check_format is True: - Class._current_filename_ = None - Class._current_instance_ = None - format_instance = Class.get_instance(filenames[0], **format_kwargs) + cls._current_filename_ = None + cls._current_instance_ = None + format_instance = cls.get_instance(filenames[0], **format_kwargs) else: format_instance = None @@ -385,7 +427,7 @@ def get_imageset( masker=None, vendor=vendor, params=params, - format=Class, + format=cls, ) ) @@ -397,7 +439,7 @@ def get_imageset( goniometer = [] scan = [] for f in filenames: - format_instance = Class(f, **format_kwargs) + format_instance = cls(f, **format_kwargs) beam.append(format_instance.get_beam()) detector.append(format_instance.get_detector()) goniometer.append(format_instance.get_goniometer()) @@ -433,7 +475,7 @@ def get_imageset( scan = format_instance.get_scan() if scan is not None: for f in filenames[1:]: - format_instance = Class(f, **format_kwargs) + format_instance = cls(f, **format_kwargs) scan += format_instance.get_scan() assert beam is not None, "Can't create Sequence without beam" @@ -454,7 +496,7 @@ def get_imageset( masker=masker, vendor=vendor, params=params, - format=Class, + format=cls, template=template, ), beam=beam, @@ -517,7 +559,7 @@ def _scan(self): long as the result is a scan.""" return None - def get_static_mask(self): + def get_static_mask(self) -> tuple[flex.int] | None: """Overload this method to override the static mask.""" return None diff --git a/src/dxtbx/format/FormatMultiImage.py b/src/dxtbx/format/FormatMultiImage.py index ee64d9ab3..e4c9d2806 100644 --- a/src/dxtbx/format/FormatMultiImage.py +++ b/src/dxtbx/format/FormatMultiImage.py @@ -1,10 +1,15 @@ from __future__ import annotations +import copy import functools import os +from pathlib import Path +from typing import Sequence, Type from scitbx.array_family import flex +import dxtbx.masking +import dxtbx.model as model from dxtbx.format.Format import Format, abstract from dxtbx.format.image import ImageBool from dxtbx.imageset import ImageSequence, ImageSet, ImageSetData, ImageSetLazy @@ -25,43 +30,53 @@ def _add_static_mask_to_iset(format_instance: Format, iset: ImageSet) -> None: class Reader: - def __init__(self, format_class, filenames, num_images=None, **kwargs): + def __init__( + self, + format_class: Type[FormatMultiImage], + filenames: Sequence[str], + num_images: int | None = None, + **kwargs, + ): self.kwargs = kwargs self.format_class = format_class assert len(filenames) == 1 self._filename = filenames[0] - if num_images is None: - format_instance = self.format_class.get_instance( - self._filename, **self.kwargs - ) - self._num_images = format_instance.get_num_images() - else: - self._num_images = num_images + self._num_images = num_images + self._format_instance = None + + def get_format_instance(self) -> FormatMultiImage: + print("Instantiating format class from Reader") + # if self._format_instance is None: + # self._format_instance = + # return self._format_instance + return self.format_class.get_instance(self._filename, **self.kwargs) def nullify_format_instance(self): self.format_class._current_instance_ = None self.format_class._current_filename_ = None + self._format_instance = None def read(self, index): - format_instance = self.format_class.get_instance(self._filename, **self.kwargs) - return format_instance.get_raw_data(index) + return self.get_format_instance().get_raw_data(index) def paths(self): return [self._filename] - def __len__(self): + def __len__(self) -> int: + if self._num_images is None: + self._num_images = self.get_format_instance().get_num_images() return self._num_images - def copy(self, filenames, num_images=None): + def copy(self, filenames: Sequence[str], num_images: int | None = None): return Reader(self.format_class, filenames, num_images, **self.kwargs) - def identifiers(self): + def identifiers(self) -> list[str]: return ["%s-%d" % (self._filename, index) for index in range(len(self))] - def is_single_file_reader(self): + def is_single_file_reader(self) -> bool: return True - def master_path(self): + def master_path(self) -> str: return self._filename @@ -101,7 +116,9 @@ def get_reader(cls): """ return functools.partial(Reader, cls) - def get_masker(self, goniometer=None): + def get_masker( + self, goniometer: model.Goniometer | None = None + ) -> dxtbx.masking.GoniometerShadowMasker | None: if ( isinstance(goniometer, MultiAxisGoniometer) and hasattr(self, "_dynamic_shadowing") @@ -115,19 +132,19 @@ def get_masker(self, goniometer=None): @classmethod def get_imageset( cls, - filenames, - beam=None, - detector=None, - goniometer=None, - scan=None, - as_sequence=False, - as_imageset=False, - single_file_indices=None, - format_kwargs=None, - template=None, - check_format=True, - lazy=False, - ): + filenames: str | Sequence[str], + beam: model.Beam | None = None, + detector: model.Detector | None = None, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + as_sequence: bool = False, + as_imageset: bool = False, + single_file_indices: Sequence[int] | None = None, + format_kwargs: dict | None = None, + template: str | None = None, + check_format: bool = True, + lazy: bool = False, + ) -> ImageSet | ImageSequence | ImageSetLazy: """ Factory method to create an imageset """ @@ -325,3 +342,190 @@ def get_imageset( _add_static_mask_to_iset(format_instance, iset) return iset + + @classmethod + def create_imagesequence( + cls, + filenames: Sequence[str], + single_file_indices: Sequence[int] | None = None, + format_kwargs: dict | None = None, + ) -> ImageSequence: + """Create an ImageSequence by reading the image file""" + return cls.get_imageset( + filenames, + single_file_indices=single_file_indices, + format_kwargs=format_kwargs, + as_sequence=True, + ) + + @classmethod + def reload_imagesequence( + cls, + filenames: Sequence[str], + single_file_indices: Sequence[int], + beam: model.Beam, + detector: model.Detector, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + format_kwargs: dict | None = None, + ) -> ImageSequence: + """ + Construct an ImageSequence from existing data + """ + # Although this accepts filenames as a sequence, for Liskov substitution + # reasons, since this is a FormatMultiImage we only support _one_ file + # that contains many images in a single dataset, not multiple files. + if len(filenames) == 0: + raise ValueError("Error: Cannot make an ImageSequence out of no files") + elif len(filenames) > 1: + raise ValueError( + "Error: FormatMultiImage.get_imagesequence only supports single multiimage files" + ) + + # Default empty kwargs + format_kwargs = format_kwargs or {} + # Make filenames absolute + filename = Path(filenames[0]).absolute() + # WIP: Ensure that we can't use this again + del filenames + + assert single_file_indices, ( + "We must have single file indices to reload ImageSequence" + ) + num_images = len(single_file_indices) + + # Get the format instance + # assert len(filenames) == 1 + # cls._current_filename_ = None + # cls._current_instance_ = None + # format_instance = cls.get_instance(filename, **format_kwargs) + + # if num_images is None: + # # As we now have the actual format class we can get the number + # # of images from here. This saves having to create another + # # format class instance in the Reader() constructor + # # NOTE: Having this information breaks internal assumptions in + # # *Lazy classes, so they have to figure this out in + # # their own time. + # num_images = format_instance.get_num_images() + + # Get some information from the format class + reader = cls.get_reader()([filename], num_images=num_images, **format_kwargs) + + # WIP: This should always be safe if we have the actual format class, + # and the assumption we have in this version is that we do + assert not cls.is_abstract() + # WIP: ... BUT, we don't have a FormatClass instance (and probably + # don't need to have one), so maybe we need to change all these + # existing examples to ClassMethod? Or, check_format=False always + # just set this blank, so maybe we never need this after initial + # load. + vendor = "" + + if not single_file_indices: + raise ValueError( + "single_file_indices can not be empty to reload FormatMultiImage ImageSequence" + ) + single_file_indices = flex.size_t(single_file_indices) + + # Check indices are sequential + if ( + max(single_file_indices) - min(single_file_indices) + == len(single_file_indices) - 1 + ): + raise ValueError("single_file_idices are not sequential") + num_images = len(single_file_indices) + + # Check the scan makes sense - we must want to use <= total images + if scan is not None: + assert scan.get_num_images() <= num_images + + # Create the masker + format_instance: FormatMultiImage = None + + loader = DeferredLoader( + cls, + filename, + format_kwargs=format_kwargs, + gonimeter=goniometer, + ) + + isetdata = ImageSetData( + reader=reader, + masker=loader.load_dynamic_mask, + vendor=vendor, + params=format_kwargs, + format=cls, + template=filename, + ) + + # Create the sequence + iset = ImageSequence( + isetdata, + beam=beam, + detector=detector, + goniometer=goniometer, + scan=scan, + indices=single_file_indices, + ) + + # Handle merging detector static mask... if necessary + assert iset.external_lookup.mask.data.empty(), ( + "Deferred static mask loading assumes nothing loaded already" + ) + + # Set up deferred instantiation and loading of the format class + iset.external_lookup.mask.set_data_generator(loader.load_static_mask) + return iset + + # WIP: Things reloading reads involving format_instance + # - [x] _add_static_mask_to_iset(format_instance, iset) + # - Added set_data_generator to ExternalLookupItem + # - [x] masker = format_instance.get_masker(goniometer=goniometer) + # - Masker is now a callable that returns the masker on request + # - [x] vendor = format_instance.get_vendortype() + # - Have just removed when reloading, for now + # - [x] reader = cls.get_reader()(filenames, num_images=num_images, **format_kwargs) + # - Reader now lazily loads rather than eagerly loads + + +class DeferredLoader: + """ + Single object to hold data related to deferred loading + + We don't want to eagerly access the raw image data on disk unless + it's actively requested. This class holds the information to do this, + and instances of it's methods passed into appropriate places so that + they can load on demand. + + Also, the objects this is going on are likely to be passed through a + pickly boundary, so needs to be statically defined. + """ + + def __init__( + self, + format: Type[Format], + filename: str, + format_kwargs: dict, + gonimeter: model.Goniometer, + ): + self.format = format + self.filename = filename + self.format_kwargs = copy.deepcopy(format_kwargs or {}) + self.goniometer = gonimeter + + def load_static_mask(self) -> ImageBool | None: + print("Deferred loading of static mask") + static_mask = self.format.get_instance( + self.filename, **self.format_kwargs + ).get_static_mask() + if static_mask: + return ImageBool(static_mask) + return None + + def load_dynamic_mask(self) -> dxtbx.masking.GoniometerShadowMasker | None: + """Lazy reading of dynamic masking""" + print("Rendering masker") + return self.format.get_instance(self.filename, **self.format_kwargs).get_masker( + goniometer=self.goniometer + ) diff --git a/src/dxtbx/format/nexus.py b/src/dxtbx/format/nexus.py index 6ea3eb7c1..72113ded6 100644 --- a/src/dxtbx/format/nexus.py +++ b/src/dxtbx/format/nexus.py @@ -1496,6 +1496,8 @@ class MaskFactory: A class to create an object to hold the pixel mask data """ + mask: tuple[flex.int] | None + def __init__(self, objects, index=None): def make_mask(dset, index): i = 0 if index is None else index diff --git a/src/dxtbx/imageset.h b/src/dxtbx/imageset.h index bc1d79b9e..9eebd19d3 100644 --- a/src/dxtbx/imageset.h +++ b/src/dxtbx/imageset.h @@ -86,6 +86,15 @@ class ExternalLookupItem { * Get the data */ Image get_data() const { + if (generator_ != boost::python::object()) { + auto generated = generator_(); + // Only replace data if we got "None" from this generator + if (generated != boost::python::object()) { + data_ = boost::python::extract>(generated)(); + } + // Discard the generator, no matter what + generator_ = boost::python::object(); + } return data_; } @@ -97,9 +106,18 @@ class ExternalLookupItem { data_ = data; } + /// Set a generator, to only load the external item data on first use + /// + /// This is a Python Callable[[], Image] function. The function will + /// be discarded after first use. + void set_data_generator(boost::python::object generator) { + generator_ = generator; + } + protected: std::string filename_; - Image data_; + mutable Image data_; + mutable boost::python::object generator_; }; /** @@ -171,12 +189,22 @@ class ImageSetData { ImageSetData(boost::python::object reader, masker_ptr masker) : reader_(reader), masker_(masker), + masker_obj_(), beams_(boost::python::len(reader)), detectors_(boost::python::len(reader)), goniometers_(boost::python::len(reader)), scans_(boost::python::len(reader)), reject_(boost::python::len(reader)) {} + ImageSetData(boost::python::object reader, boost::python::object masker) + : reader_(reader), + masker_(), + masker_obj_(masker), + beams_(boost::python::len(reader)), + detectors_(boost::python::len(reader)), + goniometers_(boost::python::len(reader)), + scans_(boost::python::len(reader)), + reject_(boost::python::len(reader)) {} /** * @returns The reader object */ @@ -188,6 +216,11 @@ class ImageSetData { * @returns The masker object */ masker_ptr masker() { + if (masker_ == nullptr && masker_obj_ != boost::python::object()) { + masker_ = boost::python::extract(masker_obj_())(); + masker_obj_ = boost::python::object(); + } + return masker_; } @@ -195,7 +228,7 @@ class ImageSetData { * @returns Does the imageset have a dynamic mask. */ bool has_dynamic_mask() const { - return masker_ != NULL; + return masker_ != nullptr || masker_obj_ != boost::python::object(); } /** @@ -440,7 +473,12 @@ class ImageSetData { std::size_t first, std::size_t last) const { DXTBX_ASSERT(last > first); - ImageSetData partial = ImageSetData(reader, masker_); + ImageSetData partial; + if (masker_ == nullptr && masker_obj_ != boost::python::object()) { + partial = ImageSetData(reader, masker_obj_); + } else { + partial = ImageSetData(reader, masker_); + } for (size_t i = 0; i < last - first; i++) { partial.beams_[i] = beams_[i + first]; partial.detectors_[i] = detectors_[i + first]; @@ -513,11 +551,18 @@ class ImageSetData { flex_type a = boost::python::extract(obj)(); // Return the image tile - return ImageTile(scitbx::af::versa >( + return ImageTile(scitbx::af::versa>( a.handle(), scitbx::af::c_grid<2>(a.accessor()))); } boost::python::object reader_; + /// Hold an object that can be called to get the masker. + /// + /// This won't be called until the masker is actually required, under + /// the assumption that accessing the masker requires accessing the raw + /// data file. + boost::python::object masker_obj_; + /// The Goniometer Masker object, if loaded (or present) std::shared_ptr masker_; scitbx::af::shared beams_; scitbx::af::shared detectors_; @@ -660,8 +705,8 @@ class ImageSet { * @returns The corrected data array */ Image get_corrected_data(std::size_t index) { - typedef scitbx::af::versa > array_type; - typedef scitbx::af::const_ref > const_ref_type; + typedef scitbx::af::versa> array_type; + typedef scitbx::af::const_ref> const_ref_type; // Get the multi-tile data, gain and pedestal DXTBX_ASSERT(index < indices_.size()); @@ -758,7 +803,7 @@ class ImageSet { std::size_t xsize = detector[i].get_image_size()[0]; std::size_t ysize = detector[i].get_image_size()[1]; scitbx::af::c_grid<2> grid(ysize, xsize); - scitbx::af::versa > data(grid, gain[i]); + scitbx::af::versa> data(grid, gain[i]); result.push_back(ImageTile(data)); } return result; @@ -797,7 +842,7 @@ class ImageSet { std::size_t xsize = detector[i].get_image_size()[0]; std::size_t ysize = detector[i].get_image_size()[1]; scitbx::af::c_grid<2> grid(ysize, xsize); - scitbx::af::versa > data(grid, pedestal[i]); + scitbx::af::versa> data(grid, pedestal[i]); result.push_back(ImageTile(data)); } return result; @@ -824,7 +869,7 @@ class ImageSet { for (std::size_t i = 0; i < detector.size(); ++i) { std::size_t xsize = detector[i].get_image_size()[0]; std::size_t ysize = detector[i].get_image_size()[1]; - mask.push_back(ImageTile(scitbx::af::versa >( + mask.push_back(ImageTile(scitbx::af::versa>( scitbx::af::c_grid<2>(ysize, xsize), true))); } return mask; @@ -854,8 +899,8 @@ class ImageSet { if (!external_mask.empty()) { DXTBX_ASSERT(external_mask.n_tiles() == mask.n_tiles()); for (std::size_t i = 0; i < mask.n_tiles(); ++i) { - scitbx::af::ref > m1 = mask.tile(i).data().ref(); - scitbx::af::const_ref > m2 = + scitbx::af::ref> m1 = mask.tile(i).data().ref(); + scitbx::af::const_ref> m2 = external_mask.tile(i).data().const_ref(); DXTBX_ASSERT(m1.accessor().all_eq(m2.accessor())); for (std::size_t j = 0; j < m1.size(); ++j) { @@ -1102,14 +1147,14 @@ class ImageSet { */ void clear_cache() { data_cache_ = DataCache(); - double_raw_data_cache_ = DataCache >(); + double_raw_data_cache_ = DataCache>(); } protected: ImageSetData data_; scitbx::af::shared indices_; DataCache data_cache_; - DataCache > double_raw_data_cache_; + DataCache> double_raw_data_cache_; Image get_raw_data_as_double(std::size_t index) { DXTBX_ASSERT(index < indices_.size()); diff --git a/src/dxtbx/imageset.py b/src/dxtbx/imageset.py index 3ace21567..793169a19 100644 --- a/src/dxtbx/imageset.py +++ b/src/dxtbx/imageset.py @@ -1,6 +1,9 @@ from __future__ import annotations +import typing +import warnings from collections.abc import Iterable +from typing import Sequence import natsort @@ -31,6 +34,12 @@ ImageSetData, ) +if typing.TYPE_CHECKING: + from dxtbx.format.Format import Format + + from . import model + + ext = boost_adaptbx.boost.python.import_ext("dxtbx_ext") __all__ = ( @@ -330,7 +339,7 @@ def get_template(self): return self.data().get_template() -def _analyse_files(filenames): +def _analyse_files(filenames: list[str]) -> list[tuple[str, list[int | None], bool]]: """Group images by filename into image sets. Params: @@ -355,13 +364,15 @@ def _indices_sequential_ge_zero(indices): return True - def _is_imageset_a_sequence(template, indices): - """Return True/False if the imageset is a sequence or not. + def _is_imageset_a_sequence(indices: Sequence[int | None]) -> bool: + """ + Check if a set of indices is sequential or not. - Where more than 1 image that follow sequential numbers are given - the images are catagorised as belonging to a sequence, otherwise they + Where more than 1 image that follow sequential numbers are given, + the images are categorised as belonging to a sequence; otherwise they belong to an image set. + A single index on it's own is never a sequence. """ if len(indices) <= 1: return False @@ -372,7 +383,7 @@ def _is_imageset_a_sequence(template, indices): file_groups = [] for template, indices in filelist_per_imageset.items(): # Check if this imageset is a sequence - is_sequence = _is_imageset_a_sequence(template, indices) + is_sequence = _is_imageset_a_sequence(indices) # Append the items to the group list file_groups.append((template, indices, is_sequence)) @@ -380,23 +391,32 @@ def _is_imageset_a_sequence(template, indices): return file_groups -# FIXME Lots of duplication in this class, need to tidy up class ImageSetFactory: """Factory to create imagesets and sequences.""" @staticmethod - def new(filenames, check_headers=False, ignore_unknown=False): - """Create an imageset or sequence - - Params: - filenames A list of filenames - check_headers Check the headers to ensure all images are valid - ignore_unknown Ignore unknown formats - - Returns: - A list of imagesets + def new( + filenames: list[str] | str, + *args, + **kwargs, + ) -> list[ImageSet | ImageSequence]: + """ + Create a list of imageset and/or sequences. + Args: + filenames: A list of filename templates. These will be expanded, and + any that appear to be a sequence will be loaded into ImageSequences. + check_headers: [OBSOLETE] Did nothing. + ignore_unknown: [OBSOLETE] Blanket ignored all errors when processing. """ + if args or kwargs: + # We used to have two parameters here - that did nothing, or hid errors. + # These were always optional, so warn the user if they are setting it. + warnings.warn( + "check_headers and ignore_unknown arguments to ImageSetFactory::new are obsolete", + DeprecationWarning, + stacklevel=2, + ) # Ensure we have enough images if isinstance(filenames, list): assert filenames @@ -406,22 +426,17 @@ def new(filenames, check_headers=False, ignore_unknown=False): raise RuntimeError("unknown argument passed to ImageSetFactory") # Analyse the filenames and group the images into imagesets. - filelist_per_imageset = _analyse_files(filenames) - + # # For each file list denoting an image set, create the imageset # and return as a list of imagesets. N.B sequences and image sets are # returned in the same list. imagesetlist = [] - for filelist in filelist_per_imageset: - try: - if filelist[2] is True: - iset = ImageSetFactory._create_sequence(filelist, check_headers) - else: - iset = ImageSetFactory._create_imageset(filelist, check_headers) - imagesetlist.append(iset) - except Exception: - if not ignore_unknown: - raise + for template, indices, is_sequence in _analyse_files(filenames): + if is_sequence: + iset = ImageSetFactory._create_sequence(template, indices) + else: + iset = ImageSetFactory._create_imageset(template, indices) + imagesetlist.append(iset) return imagesetlist @@ -445,7 +460,6 @@ def from_template( Returns: A list of sequences - """ if not check_format: assert not check_headers @@ -499,11 +513,8 @@ def from_template( return [sequence] @staticmethod - def _create_imageset(filelist, check_headers): + def _create_imageset(template: str, indices: list[int | None]): """Create an image set""" - # Extract info from filelist - template, indices, is_sequence = filelist - # Get the template format if "#" in template: filenames = _expand_template_to_sorted_filenames(template, indices) @@ -517,12 +528,13 @@ def _create_imageset(filelist, check_headers): return format_class.get_imageset(filenames, as_imageset=True) @staticmethod - def _create_sequence(filelist, check_headers): + def _create_sequence(template: str, indices: list[int | None]) -> ImageSequence: """Create a sequence""" - template, indices, is_sequence = filelist - # Expand the template if necessary if "#" in template: + assert not any(x is None for x in indices), ( + "Only accept None-indices for non-template filenames" + ) filenames = _expand_template_to_sorted_filenames(template, indices) else: filenames = [template] @@ -569,43 +581,47 @@ def make_imageset( @staticmethod def make_sequence( - template, - indices, - format_class=None, - beam=None, - detector=None, - goniometer=None, - scan=None, - check_format=True, - format_kwargs=None, - ): - """Create a sequence""" + template: str, + indices: Sequence[int], + format_class: Format | None = None, + beam: model.Beam | None = None, + detector: model.Detector | None = None, + goniometer: model.Goniometer | None = None, + scan: model.Scan | None = None, + check_format: bool = True, + format_kwargs: dict | None = None, + ) -> ImageSequence: + """ + Create a sequence + + Args: + format_class: The Format class to use. If unspecified, then + this will be either automatically detected via dxtbx (if + `check_format` is False), or default to Format or + FormatMultiImage depending on whether the template looks + like a multiimage template or not. + """ indices = sorted(indices) # Import here as Format and Imageset have cyclic dependencies from dxtbx.format.Format import Format from dxtbx.format.FormatMultiImage import FormatMultiImage - # Get the template format if "#" in template: filenames = _expand_template_to_sorted_filenames(template, indices) - # Get the format object and reader - if format_class is None: - if check_format: - format_class = dxtbx.format.Registry.get_format_class_for_file( - filenames[0] - ) - else: - format_class = Format + default_format_class = Format else: filenames = [template] - if format_class is None: - if check_format: - format_class = dxtbx.format.Registry.get_format_class_for_file( - filenames[0] - ) - else: - format_class = FormatMultiImage + default_format_class = FormatMultiImage + + if format_class is None: + if check_format: + format_class = dxtbx.format.Registry.get_format_class_for_file( + filenames[0] + ) + else: + format_class = default_format_class + assert format_class is not None # Set the image range array_range = (min(indices) - 1, max(indices)) diff --git a/src/dxtbx/masking/__init__.py b/src/dxtbx/masking/__init__.py index 4eb3d1d1f..d6f4c9c05 100644 --- a/src/dxtbx/masking/__init__.py +++ b/src/dxtbx/masking/__init__.py @@ -6,7 +6,7 @@ from scitbx import matrix from scitbx.array_family import flex -from dxtbx.model import MultiAxisGoniometer +from dxtbx.model import Goniometer, MultiAxisGoniometer try: from ..dxtbx_masking_ext import ( @@ -39,7 +39,9 @@ class GoniometerMaskerFactory: @staticmethod - def mini_kappa(goniometer, cone_opening_angle=43.60281897270362): + def mini_kappa( + goniometer: MultiAxisGoniometer, cone_opening_angle=43.60281897270362 + ) -> GoniometerShadowMasker: """Construct a GoniometerShadowMasker for a mini-kappa goniometer. This is modelled a simple cone with the opening angle specified by @@ -80,7 +82,7 @@ def mini_kappa(goniometer, cone_opening_angle=43.60281897270362): return GoniometerShadowMasker(goniometer, coords, flex.size_t(len(coords), 0)) @staticmethod - def dls_i23_kappa(goniometer): + def dls_i23_kappa(goniometer: MultiAxisGoniometer) -> GoniometerShadowMasker: """Construct a GoniometerShadowMasker for the DLS I23 Kappa goniometer. Args: @@ -183,7 +185,7 @@ def dls_i23_kappa(goniometer): return GoniometerShadowMasker(goniometer, coords, flex.size_t(len(coords), 1)) @staticmethod - def smargon(goniometer): + def smargon(goniometer: MultiAxisGoniometer) -> SmarGonShadowMasker: """Construct a SmarGonShadowMasker for the SmarGon goniometer. Args: @@ -196,7 +198,9 @@ def smargon(goniometer): return SmarGonShadowMasker(goniometer) @staticmethod - def diamond_anvil_cell(goniometer, cone_opening_angle): + def diamond_anvil_cell( + goniometer: MultiAxisGoniometer, cone_opening_angle: float + ) -> GoniometerShadowMasker: radius_height_ratio = math.tan(1 / 2 * cone_opening_angle) height = 10 # mm radius = radius_height_ratio * height diff --git a/src/dxtbx/model/__init__.py b/src/dxtbx/model/__init__.py index 095e69db7..8cbfb5efa 100644 --- a/src/dxtbx/model/__init__.py +++ b/src/dxtbx/model/__init__.py @@ -735,6 +735,11 @@ def get_template(imset): "__id__": "ImageSequence", "template": template, } + format_class = imset.get_format_class() + if not format_class.is_abstract(): + r["__format__"] = ( + f"{format_class.__module__}.{format_class.__qualname__}" + ) if imset.reader().is_single_file_reader(): r["single_file_indices"] = list(imset.indices()) elif isinstance(imset, ImageSet): @@ -950,7 +955,7 @@ def as_json( else: return text - def as_file(self, filename, **kwargs): + def as_file(self, filename: str, **kwargs): """Dump experiment list as file.""" ext = os.path.splitext(filename)[1] j_ext = [".json", ".expt"] diff --git a/src/dxtbx/model/beam.py b/src/dxtbx/model/beam.py index 3f46df22d..ca30f6438 100644 --- a/src/dxtbx/model/beam.py +++ b/src/dxtbx/model/beam.py @@ -133,19 +133,17 @@ def from_phil( return beam @staticmethod - def from_dict(dict: dict, template: dict = None) -> Beam | PolychromaticBeam: + def from_dict(d: dict, template: dict | None = None) -> Beam | PolychromaticBeam: """Convert the dictionary to a beam model""" if template is not None: - if "__id__" in dict and "__id__" in template: - assert dict["__id__"] == template["__id__"], ( + if "__id__" in d and "__id__" in template: + assert d["__id__"] == template["__id__"], ( "Beam and template dictionaries are not the same type." ) - if dict is None and template is None: - return None joint = template.copy() if template else {} - joint.update(dict) + joint.update(d) # Create the model from the joint dictionary if "probe" not in joint: diff --git a/src/dxtbx/model/experiment_list.py b/src/dxtbx/model/experiment_list.py index de193aa67..56a8d5360 100644 --- a/src/dxtbx/model/experiment_list.py +++ b/src/dxtbx/model/experiment_list.py @@ -12,7 +12,8 @@ import pickle import sys from collections.abc import Callable, Generator, Iterable -from typing import Any +from pathlib import Path +from typing import Any, AnyStr, Type, TypedDict import natsort from tqdm import tqdm @@ -21,17 +22,22 @@ from dxtbx.format.Format import Format from dxtbx.format.FormatMultiImage import FormatMultiImage from dxtbx.format.image import ImageBool, ImageDouble -from dxtbx.format.Registry import get_format_class_for_file +from dxtbx.format.Registry import get_format_class_for_file, get_format_class_index from dxtbx.imageset import ImageGrid, ImageSequence, ImageSet, ImageSetFactory from dxtbx.model import ( + Beam, BeamFactory, + Crystal, CrystalFactory, + Detector, DetectorFactory, Experiment, ExperimentList, + Goniometer, GoniometerFactory, History, ProfileModelFactory, + Scan, ScanFactory, ) from dxtbx.sequence_filenames import ( @@ -65,6 +71,16 @@ ) +class ModelDict(TypedDict): + beam: Beam | None + crystal: Crystal | None + detector: Detector | None + goniometer: Goniometer | None + profile: Any | None + scaling_model: Any | None + scan: Scan | None + + class InvalidExperimentListError(RuntimeError): """ Indicates an error whilst validating the experiment list. @@ -195,7 +211,12 @@ class ExperimentListDict: """A helper class for serializing the experiment list to dictionary (needed to save the experiment list to JSON format.""" - def __init__(self, obj, check_format=True, directory=None): + def __init__( + self, + obj: dict, + check_format=True, + directory: AnyStr | os.PathLike | None = None, + ): """Initialise. Copy the dictionary.""" # Basic check: This is a dict-like object. This can happen if e.g. we # were passed a DataBlock list instead of an ExperimentList dictionary @@ -204,7 +225,7 @@ def __init__(self, obj, check_format=True, directory=None): self._obj = copy.deepcopy(obj) self._check_format = check_format - self._directory = directory + self._directory = Path(directory) if directory else None # If this doesn't claim to be an ExperimentList, don't even try if self._obj.get("__id__") != "ExperimentList": @@ -231,7 +252,7 @@ def __init__(self, obj, check_format=True, directory=None): ) } - def _extract_models(self, name, from_dict): + def _extract_models(self, name: str, from_dict: Callable[[dict], Any]) -> Any: """ Helper function. Extract the models. @@ -273,7 +294,9 @@ def _extract_models(self, name, from_dict): if value not in mmap: mmap[value] = len(mlist) mlist.append( - from_dict(_experimentlist_from_file(value, self._directory)) + from_dict( + _experimentlist_from_file(Path(value), self._directory) + ) ) eobj[name] = mmap[value] elif not isinstance(value, int): @@ -299,10 +322,10 @@ def _load_pickle_path( the pickle file, or is None if the file is inaccessible. If there is no key entry then ("", None) is returned. """ - if param not in imageset_data: + if not imageset_data.get(param): return "", None - filename = resolve_path(imageset_data[param], directory=self._directory) + filename = resolve_path(Path(imageset_data[param]), directory=self._directory) data = None if filename: try: @@ -315,18 +338,20 @@ def _load_pickle_path( return filename, data - def _imageset_from_imageset_data(self, imageset_data, models): + def _imageset_from_imageset_data( + self, imageset_data: dict, models: ModelDict + ) -> ImageSet | ImageSequence | ImageGrid | None: """Make an imageset from imageset_data - help with refactor decode.""" assert imageset_data is not None if "params" in imageset_data: - format_kwargs = imageset_data["params"] + format_kwargs: dict = imageset_data["params"] else: format_kwargs = {} - beam = models["beam"] - detector = models["detector"] - goniometer = models["goniometer"] - scan = models["scan"] + beam: Beam = models["beam"] + detector: Detector = models["detector"] + goniometer: Goniometer | None = models["goniometer"] + scan: Scan | None = models["scan"] # Load the external lookup data mask_filename, mask = self._load_pickle_path(imageset_data, "mask") @@ -421,7 +446,7 @@ def _imageset_from_imageset_data(self, imageset_data, models): return imageset - def decode(self): + def decode(self) -> ExperimentList: """Decode the dictionary into a list of experiments.""" # Extract all the experiments - first find all scans belonging to # same imageset @@ -463,7 +488,7 @@ def decode(self): el = ExperimentList() for eobj in self._obj["experiment"]: # Get the models - identifier = eobj.get("identifier", "") + identifier: str = eobj.get("identifier", "") beam = self._lookup_model("beam", eobj) detector = self._lookup_model("detector", eobj) goniometer = self._lookup_model("goniometer", eobj) @@ -550,13 +575,13 @@ def _make_grid(self, imageset, format_kwargs=None): def _make_sequence( self, - imageset, - beam=None, - detector=None, - goniometer=None, - scan=None, - format_kwargs=None, - ): + imageset: dict, + beam: Beam | None = None, + detector: Detector | None = None, + goniometer: Goniometer | None = None, + scan: Scan | None = None, + format_kwargs: dict | None = None, + ) -> ImageSequence: """Make an image sequence.""" # Get the template format template = resolve_path(imageset["template"], directory=self._directory) @@ -568,8 +593,37 @@ def _make_sequence( else: i0, i1 = scan.get_image_range() + def _get_validate_format(name: str) -> Type[Format]: + """ + Get a Format class type instance, from the fully qualified name + + Do extra work to validate that the Format in the registry matches the + expected qualified name completely. + """ + shortname = name.split(".")[-1] + index = get_format_class_index() + if shortname not in index: + raise RuntimeError( + f"ImageSequence is registered as an unrecognised format class: '{name}'" + ) + # As a safety check, ensure that this is in the same location + concrete = index[shortname][0]() + concrete_full_name = f"{concrete.__module__}.{concrete.__qualname__}" + if concrete_full_name != name: + raise RuntimeError( + f"ImageSequence has recognised format class '{shortname}', but at an unrecognised location ({concrete} instead of the expected {name})" + ) + return concrete + format_class = None - if self._check_format is False: + + # WIP: This loads the matching format class if we have one defined. But we + # need to check to see what happens everywhere else when check_format=True + # (we want to effectively _always_ have check_format=True because the need + # to access it should have gone away). + if qualname := imageset.get("__format__"): + format_class = _get_validate_format(qualname) + elif self._check_format is False: if "single_file_indices" in imageset: format_class = FormatMultiImage @@ -586,13 +640,13 @@ def _make_sequence( format_kwargs=format_kwargs, ) - def _lookup_model(self, name, experiment_dict): + def _lookup_model(self, name: str, experiment_dict: dict[str, int]) -> dict | None: """ Find a model by looking up its index from a dictionary Args: - name (str): The model name e.g. 'beam', 'detector' - experiment_dict (Dict[str, int]): + name: The model name e.g. 'beam', 'detector' + experiment_dict: The experiment dictionary. experiment_dict[name] must exist and be not None to retrieve a model. If this key exists, then there *must* be an item with this index @@ -616,7 +670,7 @@ def _scaling_model_from_dict(obj): return entry_point.load().from_dict(obj) -def _experimentlist_from_file(filename, directory=None): +def _experimentlist_from_file(filename: Path, directory: Path | None = None) -> dict: """Load a model dictionary from a file.""" filename = resolve_path(filename, directory=directory) try: diff --git a/src/dxtbx/sequence_filenames.py b/src/dxtbx/sequence_filenames.py index 420c3f108..31b88f743 100644 --- a/src/dxtbx/sequence_filenames.py +++ b/src/dxtbx/sequence_filenames.py @@ -3,19 +3,21 @@ import os import re from collections import defaultdict +from collections.abc import Mapping from glob import glob +from typing import AnyStr import natsort -def template_regex(filename): +def template_regex(filename: str) -> tuple[str | None, int]: """Works out a template from a filename. Tries a bunch of templates to work out the most sensible. N.B. assumes that the image index will be the last digits found in the file name. - Arguments: - filename (str): The filename to template-ize + Args: + filename: The filename to template-ize Returns: Tuple[str or None, int]: @@ -115,7 +117,9 @@ def template_regex_from_list(filenames): return common_prefix + template, indices -def group_files_by_imageset(filenames): +def group_files_by_imageset( + filenames: list[AnyStr | os.PathLike], +) -> Mapping[str, list[int | None]]: """Group filenames by supposed imageset. Get the template for each file in the list. Then add to a dictionary @@ -128,7 +132,7 @@ def group_files_by_imageset(filenames): # Calculate the template for each image. If the template is None # (i.e. there are no numbers to identify the filename, add the # filename itself. - template = [] + template: list[tuple[str, int | None]] = [] for f in filenames: f = os.fspath(f) t = template_regex(f) @@ -181,7 +185,7 @@ def replace_template_format_with_hash(match): return "#" * len(match.group(0) % 0) -def template_string_to_glob_expr(template): +def template_string_to_glob_expr(template: str) -> str: """Convert the template to a glob expression.""" if template.count("#") == 1: # https://github.com/cctbx/dxtbx/issues/646 @@ -194,7 +198,7 @@ def template_string_number_index(template): return template.find("#"), template.rfind("#") + 1 -def locate_files_matching_template_string(template): +def locate_files_matching_template_string(template: str) -> list[str]: """Return all files matching template.""" matches = glob(template_string_to_glob_expr(template)) if template.count("#") != 1: @@ -206,8 +210,13 @@ def locate_files_matching_template_string(template): return [os.path.join(*m) for m in matches if patt.match(m[1])] -def template_image_range(template): - """Return the image range of files with this template.""" +def template_image_range(template: str) -> tuple[int, int]: + """ + Return the image range of files with this template. + + Raises: + ValueError: If the template doesn't match any files + """ # Find the files matching the template filenames = locate_files_matching_template_string(template) diff --git a/src/dxtbx/serialize/filename.py b/src/dxtbx/serialize/filename.py index 0f51e665e..5fe230c84 100644 --- a/src/dxtbx/serialize/filename.py +++ b/src/dxtbx/serialize/filename.py @@ -2,31 +2,37 @@ import glob import os +from pathlib import Path +from typing import TypeVar from dxtbx.sequence_filenames import template_string_to_glob_expr +T = TypeVar("T", str, Path) -def resolve_path(path, directory=None): + +def resolve_path(path: T, directory: str | os.PathLike | None = None) -> T: """Resolve a file path. First expand any environment and user variables. Then create the absolute path by applying the relative path to the provided directory, if necessary. Args: - path (str): The path to resolve - directory (Optional[str]): The local path to resolve relative links + path: The path to resolve directory: The local path to resolve relative + links Returns: - str: The absolute path to the file to read if accessible, otherwise - return the original path as provided + The absolute path to the file to read, if accessible, otherwise + return the original path. """ + + assert path, "Believe this is overly defensive, check if we ever rely on" if not path: - return "" - trial_path = os.path.expanduser(os.path.expandvars(path)) + return path + trial_path = os.path.expanduser(os.path.expandvars(str(path))) if directory and not os.path.isabs(trial_path): - trial_path = os.path.join(directory, trial_path) + trial_path = os.path.join(os.fspath(directory), trial_path) trial_path = os.path.abspath(trial_path) if glob.glob(template_string_to_glob_expr(trial_path)): - return trial_path + return type(path)(trial_path) else: - return path + return type(path)(path) diff --git a/src/dxtbx/serialize/imageset.py b/src/dxtbx/serialize/imageset.py index 6611ff7ce..e331e8a95 100644 --- a/src/dxtbx/serialize/imageset.py +++ b/src/dxtbx/serialize/imageset.py @@ -2,6 +2,7 @@ import os import pickle +from typing import AnyStr, overload from dxtbx.format.image import ImageBool, ImageDouble from dxtbx.imageset import ImageSequence, ImageSet, ImageSetFactory @@ -9,7 +10,21 @@ from dxtbx.serialize.filename import resolve_path -def filename_to_absolute(filename): +@overload +def filename_to_absolute( + filename: list[os.PathLike | AnyStr], +) -> list[AnyStr]: ... + + +@overload +def filename_to_absolute( + filename: os.PathLike | AnyStr, +) -> AnyStr: ... + + +def filename_to_absolute( + filename: list[os.PathLike | AnyStr] | os.PathLike | AnyStr, +) -> list[AnyStr] | AnyStr: """Convert filenames to absolute form.""" if isinstance(filename, list): @@ -18,7 +33,7 @@ def filename_to_absolute(filename): return os.path.abspath(filename) -def filename_or_none(filename): +def filename_or_none(filename: str | None) -> str | None: if filename is None or filename == "": return None return filename_to_absolute(filename) @@ -89,7 +104,9 @@ def imageset_to_dict(imageset): raise TypeError("Unknown ImageSet Type") -def basic_imageset_from_dict(d, directory=None): +def basic_imageset_from_dict( + d: dict, directory: str | os.PathLike | None = None +) -> ImageSet: """Construct an ImageSet class from the dictionary.""" # Get the filename list and create the imageset filenames = [resolve_path(str(p), directory=directory) for p in d["filenames"]] @@ -123,7 +140,9 @@ def basic_imageset_from_dict(d, directory=None): return imageset -def imagesequence_from_dict(d, check_format=True, directory=None): +def imagesequence_from_dict( + d: dict, check_format: bool = True, directory: AnyStr | os.PathLike | None = None +) -> ImageSequence: """Construct and image sequence from the dictionary.""" # Get the template (required) template = resolve_path(str(d["template"]), directory=directory) @@ -184,15 +203,37 @@ def imagesequence_from_dict(d, check_format=True, directory=None): return sequence -def imageset_from_dict(d, check_format=True, directory=None): - """Convert the dictionary to a sequence +@overload +def imageset_from_dict( + d: None, check_format: bool = ..., directory: str | os.PathLike | None = ... +) -> None: ... - Params: - d The dictionary of parameters + +@overload +def imageset_from_dict( + d: dict | None, + check_format: bool = ..., + directory: str | os.PathLike | None = ..., +) -> ImageSequence | ImageSet | None: ... + + +def imageset_from_dict( + d: dict | None, + check_format: bool = True, + directory: str | os.PathLike | None = None, +) -> ImageSequence | ImageSet | None: + """ + Convert a dictionary to an ImageSet + + Args: + d: The imageset model dictionary Returns: The sequence + Raises: + ValueError: If the `d["__id__"]` is not "imageset". + TypeError: If d does not contain "filename" or "template". """ # Check the input if d is None: