Add more type annotations to dxtbx for refactoring efforts

.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

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<ImageSetData::masker_ptr>(masker)();
-  }
-
   /**
    * A constructor for the imageset data class
    */
   std::shared_ptr<ImageSetData> make_imageset_data1(boost::python::object reader,
                                                     boost::python::object masker) {
-    // Create the pointer
-    std::shared_ptr<ImageSetData> 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<ImageSetData::masker_ptr> 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<ImageSetData>(self);
   }
 
   /**
@@ -70,8 +76,7 @@ namespace dxtbx { namespace boost_python {
                                                     boost::python::dict params,
                                                     boost::python::object format) {
     // Create the pointer
-    std::shared_ptr<ImageSetData> 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<T>::get_filename,
                     &ExternalLookupItem<T>::set_filename)
       .add_property(
-        "data", &ExternalLookupItem<T>::get_data, &ExternalLookupItem<T>::set_data);
+        "data", &ExternalLookupItem<T>::get_data, &ExternalLookupItem<T>::set_data)
+      .def("set_data_generator", &ExternalLookupItem<T>::set_data_generator);
   }
 
   /**

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: ...

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: ...

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