- fixing PyShape

This commit is contained in:
w.pomp
2026-07-22 16:59:04 +02:00
parent 1471ec00af
commit b063a63869
+65 -38
View File
@@ -395,6 +395,7 @@ impl PyView {
py: Python<'py>, py: Python<'py>,
n: Bound<'py, PyAny>, n: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> { ) -> PyResult<Bound<'py, PyAny>> {
// TODO: newaxis
let slice: Vec<_> = if n.is_instance_of::<PyTuple>() { let slice: Vec<_> = if n.is_instance_of::<PyTuple>() {
n.cast_into::<PyTuple>()?.into_iter().collect() n.cast_into::<PyTuple>()?.into_iter().collect()
} else if n.is_instance_of::<PyList>() { } else if n.is_instance_of::<PyList>() {
@@ -1948,51 +1949,73 @@ impl PyShape {
))] ))]
idx: Bound<'py, PyAny>, idx: Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyAny>> { ) -> PyResult<Bound<'py, PyAny>> {
let idx = if idx.is_instance_of::<PyNone>() || idx.is_instance_of::<PyEllipsis>() { let (idx, is_idx) = if idx.is_instance_of::<PyNone>() || idx.is_instance_of::<PyEllipsis>()
vec![0, 1, 2, 3, 4] {
((0..self.inner.order.len()).collect(), true)
} else if idx.is_instance_of::<PySlice>() { } else if idx.is_instance_of::<PySlice>() {
let indices = idx.cast::<PySlice>()?.indices(5)?; let indices = idx
if indices.step > 0 { .cast::<PySlice>()?
(indices.start..indices.stop) .indices(self.inner.order.len() as isize)?;
.step_by(indices.step as usize) (
.map(|i| i as usize) if indices.step > 0 {
.collect::<Vec<_>>() (indices.start..indices.stop)
} else { .step_by(indices.step as usize)
(indices.stop..indices.start) .map(|i| i as usize)
.step_by(-indices.step as usize) .collect::<Vec<_>>()
.map(|i| i as usize) } else {
.collect::<Vec<_>>() (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>() { } 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>() { } 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>() { } else if idx.is_instance_of::<PyString>() {
let s = idx.cast::<PyString>()?.extract::<String>()?; let s = idx.cast::<PyString>()?.extract::<String>()?;
s.to_uppercase() (
.chars() s.to_uppercase()
.map(|i| match i { .chars()
'C' => Ok(0), .map(|i| match i {
'Z' => Ok(1), 'C' => Ok(self.inner.c),
'T' => Ok(2), 'Z' => Ok(self.inner.z),
'Y' => Ok(3), 'T' => Ok(self.inner.t),
'X' => Ok(4), 'Y' => Ok(self.inner.y),
_ => Err(Error::Parse(s.to_string())), 'X' => Ok(self.inner.x),
}) _ => Err(Error::Parse(s.to_string())),
.collect::<Result<Vec<_>, _>>()? })
.collect::<Result<Vec<_>, _>>()?,
false,
)
} else if idx.is_instance_of::<PyInt>() { } else if idx.is_instance_of::<PyInt>() {
vec![idx.cast::<PyInt>()?.extract::<usize>()?] (vec![idx.cast::<PyInt>()?.extract::<usize>()?], true)
} else { } 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() { if shape.is_empty() {
Ok(PyNone::get(py).into_bound_py_any(py)?) Ok(PyNone::get(py).into_bound_py_any(py)?)
} else if shape.len() == 1 { } else if shape.len() == 1 {
@@ -2014,7 +2037,11 @@ impl PyShape {
#[getter] #[getter]
fn axes(&self) -> String { fn axes(&self) -> String {
self.inner.order.iter().map(|axis| format!("{}", axis)).collect::<String>() self.inner
.order
.iter()
.map(|axis| format!("{}", axis))
.collect::<String>()
} }
} }