- 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>,
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>()
}
}