pyo3_tracing_subscriber/
common.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3
4use opentelemetry::InstrumentationScope;
5/// This module contains a number of `rigetti-pyo3` ports which were
6/// backed out due to build issues involving the `pyo3/extension-module`
7/// feature. This should be replaced upon resolution of
8/// <https://github.com/rigetti/pyo3-opentelemetry/issues/15/>.
9use pyo3::prelude::*;
10use pyo3::PyErr;
11
12/// A macro for initializing a submodule.
13#[macro_export]
14macro_rules! create_init_submodule {
15    (
16        $(classes: [ $($class: ty),+ ],)?
17        $(consts: [ $($const: ident),+ ],)?
18        $(errors: [ $($error: ty),+ ],)?
19        $(funcs: [ $($func: path),+ ],)?
20        $(submodules: [ $($mod_name: literal: $init_submod: path),+ ],)?
21    ) => {
22        pub(crate) fn init_submodule(_name: &str, _py: pyo3::Python, m: &pyo3::types::PyModule) -> pyo3::PyResult<()> {
23            $($(
24            m.add_class::<$class>()?;
25            )+)?
26            $($(
27            m.add(::std::stringify!($const), $crate::ToPython::<pyo3::Py<pyo3::PyAny>>::to_python(&$const, _py)?)?;
28            )+)?
29            $($(
30            m.add(std::stringify!($error), _py.get_type::<$error>())?;
31            )+)?
32            $($(
33            m.add_function(pyo3::wrap_pyfunction!($func, m)?)?;
34            )+)?
35            $(
36                let modules = _py.import("sys")?.getattr("modules")?;
37                $(
38                let qualified_name = format!("{}.{}", _name, $mod_name);
39                let submod = pyo3::types::PyModule::new(_py, &qualified_name)?;
40                $init_submod(&qualified_name, _py, submod)?;
41                m.add($mod_name, submod)?;
42                modules.set_item(&qualified_name, submod)?;
43                )+
44            )?
45            Ok(())
46        }
47    }
48}
49
50/// A macro for wrapping a Rust error as a type error. Implements [`ToPythonError`].
51#[macro_export]
52macro_rules! py_wrap_error {
53    ($module: ident, $rust: ty, $python: ident, $base: ty) => {
54        pyo3::create_exception!($module, $python, $base);
55
56        impl $crate::common::ToPythonError for $rust {
57            fn to_py_err(self) -> pyo3::PyErr {
58                <$python>::new_err(self.to_string())
59            }
60        }
61    };
62}
63
64/// A macro for wrapping a Rust error.
65#[macro_export]
66macro_rules! wrap_error {
67    ($name: ident ($inner: ty)$(;)?) => {
68        #[derive(Debug)]
69        #[repr(transparent)]
70        pub(crate) struct $name($inner);
71
72        impl From<$inner> for $name {
73            fn from(inner: $inner) -> Self {
74                Self(inner)
75            }
76        }
77
78        impl From<$name> for $inner {
79            fn from(outer: $name) -> Self {
80                outer.0
81            }
82        }
83
84        impl ::std::fmt::Display for $name {
85            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
86                write!(f, "{}", self.0)
87            }
88        }
89
90        impl ::std::error::Error for $name {}
91    };
92}
93
94/// Converts to a Python error.
95pub(crate) trait ToPythonError {
96    /// Convert this error into a [`PyErr`](crate::pyo3::PyErr).
97    fn to_py_err(self) -> PyErr;
98}
99
100impl ToPythonError for PyErr {
101    fn to_py_err(self) -> PyErr {
102        self
103    }
104}
105
106impl ToPythonError for std::convert::Infallible {
107    fn to_py_err(self) -> PyErr {
108        unreachable!("Infallible can never happen")
109    }
110}
111
112#[pyclass(name = "InstrumentationLibrary")]
113#[derive(Debug, Clone)]
114pub(crate) struct PyInstrumentationLibrary {
115    name: String,
116    version: Option<String>,
117    schema_url: Option<String>,
118    attributes: HashMap<String, String>,
119}
120
121#[pymethods]
122impl PyInstrumentationLibrary {
123    #[new]
124    #[pyo3(signature = (name, /, version=None, schema_url=None, attributes=None))]
125    fn new(
126        name: String,
127        version: Option<String>,
128        schema_url: Option<String>,
129        attributes: Option<HashMap<String, String>>,
130    ) -> Self {
131        let attributes = attributes.unwrap_or_default();
132        Self {
133            name,
134            version,
135            schema_url,
136            attributes,
137        }
138    }
139}
140
141impl From<PyInstrumentationLibrary> for InstrumentationScope {
142    fn from(py_instrumentation_library: PyInstrumentationLibrary) -> Self {
143        let mut builder = Self::builder(Cow::from(py_instrumentation_library.name));
144        if let Some(version) = py_instrumentation_library.version {
145            builder = builder.with_version(Cow::from(version));
146        }
147        if let Some(schema_url) = py_instrumentation_library.schema_url {
148            builder = builder.with_schema_url(Cow::from(schema_url));
149        }
150        let mut attributes = Vec::new();
151        for (key, value) in py_instrumentation_library.attributes {
152            let kv = opentelemetry::KeyValue::new(
153                opentelemetry::Key::new(key),
154                opentelemetry::Value::from(value),
155            );
156            attributes.push(kv);
157        }
158        builder = builder.with_attributes(attributes);
159
160        builder.build()
161    }
162}
163
164create_init_submodule! {
165    classes: [ PyInstrumentationLibrary ],
166}