Skip to main content

rustpython_vm/function/
either.rs

1use crate::{
2    AsObject, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine, convert::ToPyObject,
3};
4use core::borrow::Borrow;
5
6pub enum Either<A, B> {
7    A(A),
8    B(B),
9}
10
11impl<A, B> Borrow<PyObject> for Either<A, B>
12where
13    A: Borrow<PyObject>,
14    B: Borrow<PyObject>,
15{
16    #[inline(always)]
17    fn borrow(&self) -> &PyObject {
18        match self {
19            Self::A(a) => a.borrow(),
20            Self::B(b) => b.borrow(),
21        }
22    }
23}
24
25impl<A, B> AsRef<PyObject> for Either<A, B>
26where
27    A: AsRef<PyObject>,
28    B: AsRef<PyObject>,
29{
30    #[inline(always)]
31    fn as_ref(&self) -> &PyObject {
32        match self {
33            Self::A(a) => a.as_ref(),
34            Self::B(b) => b.as_ref(),
35        }
36    }
37}
38
39impl<A, B> From<Either<A, B>> for PyObjectRef
40where
41    A: Into<Self>,
42    B: Into<Self>,
43{
44    #[inline(always)]
45    fn from(value: Either<A, B>) -> Self {
46        match value {
47            Either::A(a) => a.into(),
48            Either::B(b) => b.into(),
49        }
50    }
51}
52
53impl<A, B> ToPyObject for Either<A, B>
54where
55    A: ToPyObject,
56    B: ToPyObject,
57{
58    #[inline(always)]
59    fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef {
60        match self {
61            Self::A(a) => a.to_pyobject(vm),
62            Self::B(b) => b.to_pyobject(vm),
63        }
64    }
65}
66
67/// This allows a builtin method to accept arguments that may be one of two
68/// types, raising a `TypeError` if it is neither.
69///
70/// # Example
71///
72/// ```
73/// use rustpython_vm::VirtualMachine;
74/// use rustpython_vm::builtins::{PyStrRef, PyIntRef};
75/// use rustpython_vm::function::Either;
76///
77/// fn do_something(arg: Either<PyIntRef, PyStrRef>, vm: &VirtualMachine) {
78///     match arg {
79///         Either::A(int)=> {
80///             // do something with int
81///         }
82///         Either::B(string) => {
83///             // do something with string
84///         }
85///     }
86/// }
87/// ```
88impl<A, B> TryFromObject for Either<A, B>
89where
90    A: TryFromObject,
91    B: TryFromObject,
92{
93    fn try_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {
94        A::try_from_object(vm, obj.clone())
95            .map(Either::A)
96            .or_else(|_| B::try_from_object(vm, obj.clone()).map(Either::B))
97            .map_err(|_| vm.new_type_error(format!("unexpected type {}", obj.class())))
98    }
99}