qcs_sdk/
client.rs

1use std::process::Output;
2
3use pyo3::{exceptions::PyValueError, pyfunction, PyAny};
4use qcs_api_client_common::configuration::{
5    AuthServer, ClientConfigurationBuilder, ClientConfigurationBuilderError, ClientCredentials,
6    ExternallyManaged, OAuthSession, RefreshToken,
7};
8use rigetti_pyo3::{
9    create_init_submodule, py_async, py_function_sync_async, py_wrap_error, py_wrap_type,
10    pyo3::{
11        conversion::IntoPy, exceptions::PyRuntimeError, pyclass::CompareOp, pymethods, PyObject,
12        PyResult, Python,
13    },
14    wrap_error, PyWrapper, ToPythonError,
15};
16
17use qcs::client::{self, Qcs};
18use tokio_util::sync::CancellationToken;
19
20use crate::py_sync;
21
22create_init_submodule! {
23    classes: [
24        PyQcsClient,
25        OAuthSession,
26        AuthServer,
27        RefreshToken,
28        ClientCredentials,
29        ExternallyManaged
30    ],
31    errors: [
32        LoadClientError,
33        BuildClientError
34    ],
35}
36
37wrap_error!(RustLoadClientError(client::LoadError));
38py_wrap_error!(client, RustLoadClientError, LoadClientError, PyRuntimeError);
39
40wrap_error!(RustBuildClientError(ClientConfigurationBuilderError));
41
42py_wrap_error!(
43    client,
44    RustBuildClientError,
45    BuildClientError,
46    PyRuntimeError
47);
48
49py_wrap_type! {
50    PyQcsClient(Qcs) as "QCSClient";
51}
52
53impl PyQcsClient {
54    pub fn get_or_create_client(client: Option<Self>) -> Qcs {
55        match client {
56            Some(client) => client.into(),
57            None => Qcs::load(),
58        }
59    }
60}
61
62impl PartialEq for PyQcsClient {
63    fn eq(&self, other: &Self) -> bool {
64        format!("{:?}", self.0) == format!("{:?}", other.0)
65    }
66}
67
68#[pymethods]
69impl PyQcsClient {
70    #[new]
71    #[pyo3(signature = (
72        /,
73        oauth_session = None,
74        api_url = None,
75        grpc_api_url = None,
76        quilc_url = None,
77        qvm_url = None
78    ))]
79    pub fn new(
80        oauth_session: Option<OAuthSession>,
81        api_url: Option<String>,
82        grpc_api_url: Option<String>,
83        quilc_url: Option<String>,
84        qvm_url: Option<String>,
85    ) -> PyResult<Self> {
86        let mut builder = ClientConfigurationBuilder::default();
87        if let Some(session) = oauth_session {
88            builder.oauth_session(Some(session));
89        }
90        if let Some(api_url) = api_url {
91            builder.api_url(api_url);
92        }
93        if let Some(grpc_api_url) = grpc_api_url {
94            builder.grpc_api_url(grpc_api_url);
95        }
96        if let Some(quilc_url) = quilc_url {
97            builder.quilc_url(quilc_url);
98        }
99        if let Some(qvm_url) = qvm_url {
100            builder.qvm_url(qvm_url);
101        }
102        let client = builder
103            .build()
104            .map(Qcs::with_config)
105            .map_err(RustBuildClientError::from)
106            .map_err(RustBuildClientError::to_py_err)?;
107
108        Ok(Self(client))
109    }
110
111    #[staticmethod]
112    #[pyo3(signature = (/, profile_name = None))]
113    fn load(profile_name: Option<String>) -> PyResult<Self> {
114        Ok(match profile_name {
115            Some(profile_name) => Qcs::with_profile(profile_name)
116                .map(PyQcsClient)
117                .map_err(RustLoadClientError)
118                .map_err(RustLoadClientError::to_py_err)?,
119            None => Self(Qcs::load()),
120        })
121    }
122
123    #[staticmethod]
124    #[pyo3(signature = (/, profile_name = None))]
125    fn load_with_login(py: Python<'_>, profile_name: Option<String>) -> PyResult<Self> {
126        do_until_ctrl_c(move |cancel_token| {
127            py_sync!(py, async move {
128                Qcs::with_login(cancel_token, profile_name)
129                    .await
130                    .map(PyQcsClient)
131                    .map_err(RustLoadClientError::from)
132                    .map_err(RustLoadClientError::to_py_err)
133            })
134        })
135    }
136
137    #[staticmethod]
138    #[pyo3(signature = (/, profile_name = None))]
139    fn load_with_login_async(py: Python<'_>, profile_name: Option<String>) -> PyResult<&PyAny> {
140        do_until_ctrl_c(move |cancel_token| {
141            py_async!(py, async move {
142                Qcs::with_login(cancel_token, profile_name)
143                    .await
144                    .map(PyQcsClient)
145                    .map_err(RustLoadClientError::from)
146                    .map_err(RustLoadClientError::to_py_err)
147            })
148        })
149    }
150
151    #[getter]
152    pub fn api_url(&self) -> String {
153        self.as_ref().get_config().api_url().to_string()
154    }
155
156    #[getter]
157    pub fn grpc_api_url(&self) -> String {
158        self.as_ref().get_config().grpc_api_url().to_string()
159    }
160
161    #[getter]
162    pub fn quilc_url(&self) -> String {
163        self.as_ref().get_config().quilc_url().to_string()
164    }
165
166    #[getter]
167    pub fn qvm_url(&self) -> String {
168        self.as_ref().get_config().qvm_url().to_string()
169    }
170
171    #[getter]
172    pub fn oauth_session(&self, py: Python<'_>) -> PyResult<OAuthSession> {
173        py_get_oauth_session(py, self.clone())
174    }
175
176    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyObject {
177        match op {
178            CompareOp::Eq => (self == other).into_py(py),
179            CompareOp::Ne => (self != other).into_py(py),
180            _ => py.NotImplemented(),
181        }
182    }
183}
184
185py_function_sync_async! {
186    #[pyfunction]
187    async fn get_oauth_session(client: PyQcsClient) -> PyResult<OAuthSession> {
188        client.as_inner().get_config().oauth_session().await.map_err(|e| PyValueError::new_err(e.to_string()))
189
190    }
191}
192
193/// Run the given function with a CancellationToken that is cancelled when `Ctrl+C` is pressed.
194fn do_until_ctrl_c<T>(f: impl FnOnce(CancellationToken) -> T) -> T {
195    let cancel_token = CancellationToken::new();
196    let cancel_token_ctrl_c = cancel_token.clone();
197    tokio::spawn(cancel_token.clone().run_until_cancelled_owned(async move {
198        let _ = tokio::signal::ctrl_c().await;
199        cancel_token_ctrl_c.cancel();
200    }));
201
202    let value = f(cancel_token.clone());
203    cancel_token.cancel();
204    value
205}