Accepting Union types for function arguments

[mehcode]

I'd like to define a #[pyfunction] that accepts either a String or some custom #[pyclass].

[davidhewitt]

Renaming to mention Python's Union type - which is what this issue is really talking about integration for imo.

[birkenfeld]

I think this is a hopelessly complex topic and you're better off taking a &PyAny and doing the extraction yourself.

(Stdlib functions implemented C do the same.)

[davidhewitt]

I agree that this is a complete moonshot and unlikely to be supported for a long long time! If rust had variadic generics a la C++, this would be easy ;)

[sebpuetz]

You can probably emulate this behaviour to the point where it's fairly similar to Python where you'd often need something like isinstance. The if instance(obj, type) in Python would be replaced by a match in Rust.

Say you want a method to take Union[str, int], then you could implement some wrapper enum that you implement extraction logic for:

use pyo3::prelude::*;
use pyo3::types::{PyString, PyLong};
use pyo3::{PyTypeInfo, exceptions, wrap_pyfunction};

#[pyfunction]
pub fn bar(foo: Union) {
    match foo {
        Union::Str(s) => println!("{}", s),
        Union::Int(i) => println!("{}", i),
    }
}

pub enum Union {
    Str(String),
    Int(isize),
}

impl<'source> FromPyObject<'source> for Union {
    fn extract(ob: &'source PyAny) -> PyResult<Self> {
        if let Ok(s) = PyString::try_from(ob) {
            return Ok(Union::Str(s.to_string()))
        } else if let Ok(i) = PyLong::try_from(ob) {
            return Ok(Union::Int(i.extract()?))
        } else {
            Err(exceptions::PyTypeError::py_err(format!("Can't convert from {} to Union[str, int]", ob.get_type().name())))
        }

    }
}

#[pymodule]
fn foo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(bar))?;
    Ok(())
}

For Python there's essentially no difference to using a Python def'd function dealing with both int and str.

>>> import foo
>>> foo.bar("Test")
Test
>>> foo.bar(1)
1
>>> foo.bar(None)
TypeError: Can't convert from NoneType to Union[str, int]

E.g. in pure Python something like this would probably be used:

def union_fn(inp: Union[str, int]):
    if isinstance(inp, str):
        return do_stringly_things(inp)
    if isinstance(inp, int):
        return do_inty_things(inp)
    raise TypeError(f"Expected Union[str, int], got {type(inp).__name__}")

[birkenfeld]

It can be even simpler in extract with something like

if let Ok(s) = ob.extract() {
    return Ok(Union::Str(s))
} else if let Ok(i) = ob.extract() {
    ...

I guess a derive macro for FromPyObject for such enums could be in the scope of pyo3.

[sebpuetz]

Iiirc, FromPyObject::extract introduces some overhead for failed conversions since it creates a PyErr which includes the formatted underlying TypeErrorr and even entails acquiring the GIL.

So I think the way to go about this is either use PyTryFrom and accept that failing extractions up to the correct type create errors a PyDowncastError or take the type check and downcast out of the trait implementation:

        unsafe {
            if T::is_instance(value) {
                Ok(Self::try_from_unchecked(value))
            }
        }

Perhaps also add an option to the potential proc macro that allows is_instance and is_exact_instance to be used.

---

While trying to figure out what's possible, I actually wrote an wrapper enum for the native types:

pub enum PyNative<'a> {
    // PySequence and PyIter are excluded here
    Object(&'a PyAny),
    None(&'a PyAny),
    String(&'a PyString),
    Bool(&'a PyBool),
    ByteArray(&'a PyByteArray),
    Bytes(&'a PyBytes),
    Complex(&'a PyComplex),
    DateTime(&'a PyDateTime),
    Dict(&'a PyDict),
    Float(&'a PyFloat),
    List(&'a PyList),
    Long(&'a PyLong),
    Module(&'a PyModule),
    Set(&'a PySet),
    Slice(&'a PySlice),
    Tuple(&'a PyTuple),
    Type(&'a PyType),
}

macro_rules! impl_extract_native {
    ($into_ty:ty, $variant:path, $ob:expr) => {
        unsafe {
            if <$into_ty as PyTypeInfo>::is_instance($ob) {
                return Ok($variant(<$into_ty>::try_from_unchecked($ob)))
            }
        }
    }
}

impl<'source> FromPyObject<'source> for PyNative<'source> {
    fn extract(ob: &'source PyAny) -> PyResult<Self> {
        // XXX PyNative::None instead of None(ob) ?
        if ob.is_none() {
            return Ok(PyNative::None(ob))
        }
        use PyNative::*;
        impl_extract_native!(PyBool, Bool, ob);
        impl_extract_native!(PyByteArray, ByteArray, ob);
        impl_extract_native!(PyBytes, Bytes, ob);
        impl_extract_native!(PyComplex, Complex, ob);
        impl_extract_native!(PyDateTime, DateTime, ob);
        impl_extract_native!(PyDict, Dict, ob);
        impl_extract_native!(PyFloat, Float, ob);
        impl_extract_native!(PyList, List, ob);
        impl_extract_native!(PyModule, Module, ob);
        impl_extract_native!(PyLong, Long, ob);
        impl_extract_native!(PySet, Set, ob);
        impl_extract_native!(PySlice, Slice, ob);
        impl_extract_native!(PyString, String, ob);
        impl_extract_native!(PyTuple, Tuple, ob);
        impl_extract_native!(PyType, Type, ob);
        Ok(PyNative::Object(ob))
    }
}

This is most likely not usable as-is, but if we can derive a FromPyObject impl like this for wrapper enums, it'd be a great QoL improvement.

[sebpuetz]

This can probably be closed via #1065

[davidhewitt]

I'd like to just add a final note to docs later on the conversion table.