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
193fn 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}