From b063a638696a579eafe418bd2f46bea0f670fba3 Mon Sep 17 00:00:00 2001
From: "w.pomp" <w.pomp@nki.nl>
Date: Wed, 22 Jul 2026 16:59:04 +0200
Subject: [PATCH] - fixing PyShape

---
 src/py.rs | 103 ++++++++++++++++++++++++++++++++++--------------------
 1 file changed, 65 insertions(+), 38 deletions(-)

diff --git a/src/py.rs b/src/py.rs
index 8610858..e9bd460 100644
--- a/src/py.rs
+++ b/src/py.rs
@@ -395,6 +395,7 @@ impl PyView {
         py: Python<'py>,
         n: Bound<'py, PyAny>,
     ) -> PyResult<Bound<'py, PyAny>> {
+        // TODO: newaxis
         let slice: Vec<_> = if n.is_instance_of::<PyTuple>() {
             n.cast_into::<PyTuple>()?.into_iter().collect()
         } else if n.is_instance_of::<PyList>() {
@@ -1948,51 +1949,73 @@ impl PyShape {
         ))]
         idx: Bound<'py, PyAny>,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let idx = if idx.is_instance_of::<PyNone>() || idx.is_instance_of::<PyEllipsis>() {
-            vec![0, 1, 2, 3, 4]
+        let (idx, is_idx) = if idx.is_instance_of::<PyNone>() || idx.is_instance_of::<PyEllipsis>()
+        {
+            ((0..self.inner.order.len()).collect(), true)
         } else if idx.is_instance_of::<PySlice>() {
-            let indices = idx.cast::<PySlice>()?.indices(5)?;
-            if indices.step > 0 {
-                (indices.start..indices.stop)
-                    .step_by(indices.step as usize)
-                    .map(|i| i as usize)
-                    .collect::<Vec<_>>()
-            } else {
-                (indices.stop..indices.start)
-                    .step_by(-indices.step as usize)
-                    .map(|i| i as usize)
-                    .collect::<Vec<_>>()
-            }
+            let indices = idx
+                .cast::<PySlice>()?
+                .indices(self.inner.order.len() as isize)?;
+            (
+                if indices.step > 0 {
+                    (indices.start..indices.stop)
+                        .step_by(indices.step as usize)
+                        .map(|i| i as usize)
+                        .collect::<Vec<_>>()
+                } else {
+                    (indices.stop..indices.start)
+                        .step_by(-indices.step as usize)
+                        .map(|i| i as usize)
+                        .collect::<Vec<_>>()
+                },
+                true,
+            )
         } else if idx.is_instance_of::<PyList>() {
-            idx.cast::<PyList>()?.extract::<Vec<usize>>()?
+            (idx.cast::<PyList>()?.extract::<Vec<usize>>()?, true)
         } else if idx.is_instance_of::<PyTuple>() {
-            idx.cast::<PyTuple>()?.extract::<Vec<usize>>()?
+            (idx.cast::<PyTuple>()?.extract::<Vec<usize>>()?, true)
         } else if idx.is_instance_of::<PyString>() {
             let s = idx.cast::<PyString>()?.extract::<String>()?;
-            s.to_uppercase()
-                .chars()
-                .map(|i| match i {
-                    'C' => Ok(0),
-                    'Z' => Ok(1),
-                    'T' => Ok(2),
-                    'Y' => Ok(3),
-                    'X' => Ok(4),
-                    _ => Err(Error::Parse(s.to_string())),
-                })
-                .collect::<Result<Vec<_>, _>>()?
+            (
+                s.to_uppercase()
+                    .chars()
+                    .map(|i| match i {
+                        'C' => Ok(self.inner.c),
+                        'Z' => Ok(self.inner.z),
+                        'T' => Ok(self.inner.t),
+                        'Y' => Ok(self.inner.y),
+                        'X' => Ok(self.inner.x),
+                        _ => Err(Error::Parse(s.to_string())),
+                    })
+                    .collect::<Result<Vec<_>, _>>()?,
+                false,
+            )
         } else if idx.is_instance_of::<PyInt>() {
-            vec![idx.cast::<PyInt>()?.extract::<usize>()?]
+            (vec![idx.cast::<PyInt>()?.extract::<usize>()?], true)
         } else {
-            return Err(PyErr::new::<PyTypeError, _>("Unknown type"));
+            return Err(PyErr::new::<PyTypeError, _>(format!(
+                "Unknown type: {:?}",
+                idx
+            )));
+        };
+        let shape = if is_idx {
+            let mut shape = Vec::new();
+            for axis in &self.inner.order {
+                match axis {
+                    Axis::C => shape.push(self.inner.c),
+                    Axis::Z => shape.push(self.inner.z),
+                    Axis::T => shape.push(self.inner.t),
+                    Axis::Y => shape.push(self.inner.y),
+                    Axis::X => shape.push(self.inner.x),
+                    Axis::New => shape.push(1),
+                }
+            }
+            idx.into_iter()
+                .map(|i| shape[i % shape.len()])
+                .collect::<Vec<_>>()
+        } else {
+            idx
         };
-        let shape = [
-            self.inner.c,
-            self.inner.z,
-            self.inner.t,
-            self.inner.y,
-            self.inner.x,
-        ];
-        let shape = idx.into_iter().map(|i| shape[i % 5]).collect::<Vec<_>>();
         if shape.is_empty() {
             Ok(PyNone::get(py).into_bound_py_any(py)?)
         } else if shape.len() == 1 {
@@ -2014,7 +2037,11 @@ impl PyShape {
 
     #[getter]
     fn axes(&self) -> String {
-        self.inner.order.iter().map(|axis| format!("{}", axis)).collect::<String>()
+        self.inner
+            .order
+            .iter()
+            .map(|axis| format!("{}", axis))
+            .collect::<String>()
     }
 }