quil-rs 0.37.0

//! This binary is used to generate Python stub files (type hints) for the `quil` package.
//! For more information on why this exists as a separate binary rather than a build script,
//! see the [`pyo3-stub-gen`][] documentation.
//!
//! [`pyo3-stub-gen`]: https://github.com/Jij-Inc/pyo3-stub-gen

#[cfg(not(feature = "stubs"))]
mod main {
    use std::process::ExitCode;

    pub fn main() -> ExitCode {
        eprintln!("Executing this binary only makes sense with the --stubs feature enabled.");
        ExitCode::FAILURE
    }
}
#[cfg(feature = "stubs")]
/// Our stub generation code generates the `.pyi` files with
/// [`py-stub-gen`](https://github.com/Jij-Inc/pyo3-stub-gen), and then applies four categories of
/// edits to the generated `.pyi` files in order to work around the limitations of
/// [PyO3](https://pyo3.rs/) and `py-stub-gen` when it comes to generic Python types.  The syntax
/// used for generics is the old-style, pre-3.12 syntax, since we support Python 3.10 and 3.11.
///
/// We allow specifying four kinds of edits.  Each edit operates on a single line, but they are
/// shown on multiple lines here for clarity.
///
/// 1. Classes may be altered to have generic parameters.  For instance, this can replace `class
///    Waveform:` with `class Waveform(typing.Generic[_Real, _Complex]):`.  (In the new syntax,
///    where we also wouldn't have to use underscores, that would be `class Waveform[Real, Complex =
///    Real]:`.)  This currently does not support classes that have explicit superclass lists.
///
/// 2. Methods may be altered to have generic parameters.  This is a no-op, as we use the old Python
///    syntax with globally-defined `TypeVar`s; nevertheless, the information about the type
///    variables is used to generate the list of `TypeVar`s.  In the new syntax, where we also
///    wouldn't have to use underscores, this would replace
///
///    ```text
///    def evaluate(
///        self,
///        real: collections.abc.Callable[[Real], OtherReal],
///        complex: collections.abc.Callable[[Complex], OtherComplex]
///    ) -> Waveform[OtherReal, OtherComplex]:
///        ...
///    ```
///    
///    with
///    
///    ```text
///    def evaluate[OtherReal, OtherComplex = OtherReal](
///        self,
///        real: collections.abc.Callable[[Real], OtherReal],
///        complex: collections.abc.Callable[[Complex], OtherComplex]
///    ) -> Waveform[OtherReal, OtherComplex]:
///    ```
///
/// 3. Methods may have a different type annotation placed on `self`.  For instance, this can
///    replace
///
///    ```text
///    def iq_values_at_sample_rate(
///        self,
///        common: CommonBuiltinParameters[builtins.float, __T],
///        sample_rate: builtins.float
///    ) -> IqSamples:
///        ...
///    ```
///    
///    with
///    
///    ```text
///    def iq_values_at_sample_rate(
///        self: BuiltinWaveform[builtins.float, builtins.complex],
///        common: CommonBuiltinParameters[builtins.float, __T],
///        sample_rate: builtins.float
///    ) -> IqSamples:
///        ...
///    ```
///
/// 4. Anywhere `$SELF` occurs in a class that is being edited, it is replaced with the
///    *fully-parameterized* form of the class.  This is important for code generation from macros.
///    For instance, this can replace
///
///    ```text
///    def __new__(
///        cls,
///        *,
///        duration: builtins.float,
///        scale: typing.Optional[_Real] = None,
///        phase: typing.Optional[_Real] = None,
///        detuning: typing.Optional[_Real] = None
///    ) -> $SELF:
///        ...
///    ```
///
///    with
///
///    ```text
///    def __new__(
///        cls,
///        *,
///        duration: builtins.float,
///        scale: typing.Optional[_Real] = None,
///        phase: typing.Optional[_Real] = None,
///        detuning: typing.Optional[_Real] = None
///    ) -> CommonBuiltinParameters[_Real, _Complex]:
///        ...
///    ```
///
/// In addition to these edits, the tool will start the edited `.pyi` file with a list of `TypeVar`
/// definitions.  These definitions support type parameter defaults in versions of Python greater
/// than or equal to 3.13 (the earliest this became possible for explicit `TypeVar`s).  For
/// instance, this might look like:
///
/// ```text
/// import sys
/// if sys.version_info >= (3, 13):
///     _Real = typing.TypeVar("_Real")
///     _Complex = typing.TypeVar("_Complex", default=Real)
///     __T = typing.TypeVar("__T")
/// else:
///     _Real = typing.TypeVar("_Real")
///     _Complex = typing.TypeVar("_Complex")
///     __T = typing.TypeVar("__T")
///     ```
///
/// This means that **any type variables that aren't defined at runtime need to be prefixed with an
/// underscore**.  Alas!
mod main {
    use indexmap::IndexMap;

    mod edits {
        use std::fmt;

        use anyhow::bail;
        use indexmap::IndexMap;

        pub type Modules = IndexMap<String, IndexMap<String, Class>>;

        #[derive(Clone, Debug)]
        pub struct Class {
            pub type_parameters: Vec<TypeParameter>,
            pub methods: IndexMap<String, Method>,
        }

        #[derive(Clone, Debug)]
        pub struct Method {
            pub type_parameters: Vec<TypeParameter>,
            pub self_type: Option<String>,
        }

        #[derive(Clone, PartialEq, Eq, Hash, Debug)]
        pub struct TypeParameter {
            pub name: String,
            pub default: Option<String>,
        }

        impl fmt::Display for TypeParameter {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let Self { name, default } = self;
                write!(f, "{name}")?;
                if let Some(default) = default {
                    write!(f, " = {default}")?;
                }
                Ok(())
            }
        }

        pub fn collect_all_type_parameters<'a>(
            module: &'_ str,
            classes: &'a IndexMap<String, Class>,
        ) -> anyhow::Result<IndexMap<&'a str, Option<&'a str>>> {
            let mut type_parameters_with_defaults = IndexMap::<&'a str, Option<&'a str>>::new();

            let mut insert_all = |type_parameters: &'a [TypeParameter]| {
                for TypeParameter { name, default } in type_parameters {
                    match type_parameters_with_defaults.entry(name) {
                        indexmap::map::Entry::Occupied(entry) => {
                            let existing_default = *entry.get();
                            if default.as_deref() != existing_default {
                                bail!(
                                    "conflicting defaults in {module} for type variable {name}: \
                                 found both {existing_default:?} and {default:?}"
                                )
                            }
                        }
                        indexmap::map::Entry::Vacant(entry) => {
                            entry.insert(default.as_deref());
                        }
                    }
                }
                Ok(())
            };

            for Class {
                type_parameters,
                methods,
            } in classes.values()
            {
                insert_all(type_parameters)?;

                for Method {
                    type_parameters,
                    self_type: _,
                } in methods.values()
                {
                    insert_all(type_parameters)?;
                }
            }

            Ok(type_parameters_with_defaults)
        }

        impl TypeParameter {
            pub fn new<S: Into<String>>(name: S) -> Self {
                Self {
                    name: name.into(),
                    default: None,
                }
            }

            pub fn defaulted<S1: Into<String>, S2: Into<String>>(name: S1, default: S2) -> Self {
                Self {
                    name: name.into(),
                    default: Some(default.into()),
                }
            }
        }

        pub fn fmt_type_parameters_maybe_default(
            parameters: &[TypeParameter],
            defaults: bool,
        ) -> impl fmt::Display + use<'_> {
            fmt::from_fn(move |f| {
                if parameters.is_empty() {
                    return Ok(());
                }
                write!(f, "[")?;
                let mut first = true;
                for p in parameters {
                    if first {
                        first = false;
                    } else {
                        write!(f, ", ")?;
                    }
                    if defaults {
                        write!(f, "{p}")?;
                    } else {
                        write!(f, "{}", p.name)?;
                    }
                }
                write!(f, "]")
            })
        }

        pub fn fmt_type_parameters_defaultless(
            parameters: &[TypeParameter],
        ) -> impl fmt::Display + use<'_> {
            fmt_type_parameters_maybe_default(parameters, false)
        }

        fn fmt_type_var_definitions_indented_maybe_default<'a>(
            parameters: &'a IndexMap<&str, Option<&str>>,
            indentation: &'a str,
            defaults: bool,
        ) -> impl fmt::Display + use<'a> {
            fmt::from_fn(move |f| {
                for (name, default) in parameters {
                    write!(f, r#"{indentation}{name} = typing.TypeVar("{name}""#)?;
                    if defaults {
                        if let Some(default) = default {
                            write!(f, r", default = {default}")?;
                        }
                    }
                    writeln!(f, ")")?;
                }
                Ok(())
            })
        }

        pub fn fmt_type_var_definitions<'a>(
            parameters: &'a IndexMap<&str, Option<&str>>,
        ) -> impl fmt::Display + use<'a> {
            fmt::from_fn(move |f| {
                if parameters.is_empty() {
                    return Ok(());
                }

                writeln!(f, "import sys")?;
                writeln!(f, "if sys.version_info >= (3, 13):")?;
                write!(
                    f,
                    "{}",
                    fmt_type_var_definitions_indented_maybe_default(parameters, "    ", true)
                )?;
                writeln!(f, "else:")?;
                write!(
                    f,
                    "{}",
                    fmt_type_var_definitions_indented_maybe_default(parameters, "    ", false)
                )?;
                writeln!(f)
            })
        }

        pub fn fmt_class_type_parameters_as_inheritance(
            parameters: &[TypeParameter],
        ) -> impl fmt::Display + use<'_> {
            fmt::from_fn(move |f| {
                if parameters.is_empty() {
                    return Ok(());
                }
                write!(
                    f,
                    "(typing.Generic{})",
                    fmt_type_parameters_defaultless(parameters)
                )
            })
        }
    }

    mod editor {
        use std::{
            fs::File,
            io::{self, BufRead, BufReader, BufWriter, Write as _},
            path::{Path, PathBuf},
        };

        use anyhow::bail;
        use indexmap::{IndexMap, IndexSet};

        use super::edits;

        /// Try to open a file.  Succeeds with `Some((path, file))` if the file was openable, suceeds
        /// with `None` if the file couldn't be opened because it doesn't exist, and fails in all other
        /// cases.
        fn try_open_file<P: AsRef<Path>>(path: P) -> io::Result<Option<(P, BufReader<File>)>> {
            match File::open(path.as_ref()) {
                Ok(file) => Ok(Some((path, BufReader::new(file)))),
                Err(err) => {
                    if err.kind() == io::ErrorKind::NotFound {
                        Ok(None)
                    } else {
                        Err(err)
                    }
                }
            }
        }

        /// Given a path to the `root` where the Python `.pyi` files are stored and a Python `module`
        /// name, attempt to open the `.pyi` file for that module.
        fn open_python_module(
            root: &Path,
            module: &str,
        ) -> anyhow::Result<(PathBuf, BufReader<File>)> {
            let module_path = root.join(module.replace('.', "/"));

            if let Some(success) = try_open_file(module_path.with_added_extension(".pyi"))? {
                return Ok(success);
            }

            if let Some(success) = try_open_file(module_path.join("__init__.pyi"))? {
                return Ok(success);
            }

            bail!("no stub file found for {module}")
        }

        /// An editor for adjusting `.pyi` files.
        #[derive(Debug)]
        struct PyiEditor<'a> {
            context: PyiContext<'a>,
            state: PyiEditorState<'a>,
            input: FileInput,
            output: FileOutput,
        }

        /// Global, unchanging information about a `.pyi` file.
        #[derive(Debug)]
        struct PyiContext<'a> {
            module_name: &'a str,
            classes: &'a IndexMap<String, edits::Class>,
        }

        /// Information about the current class a [`PyiEditor`] is looking at, as it goes through the
        /// file line by line.
        #[derive(Debug)]
        struct PyiCurrentClass<'a> {
            class_name: &'a str,
            class: &'a edits::Class,
            unseen_methods: IndexSet<&'a str>,
        }

        /// The updatable state of the editor.
        #[derive(Debug)]
        struct PyiEditorState<'a> {
            wrote_type_vars: bool,
            unseen_classes: IndexSet<&'a str>,
            current_class: Option<PyiCurrentClass<'a>>,
        }

        /// Information about an input file.
        #[derive(Debug)]
        struct FileInput {
            input_path: PathBuf,
            input: BufReader<File>,
        }

        /// Information about an output file.
        #[derive(Debug)]
        struct FileOutput {
            output_path: PathBuf,
            output: BufWriter<File>,
        }

        impl<'a> PyiEditor<'a> {
            /// Given a path to the `root` where the Python `.pyi` files are stored, a Python
            /// `module` name, and a map of edits to be made to `classes` in that module, construct
            /// the editor that will perform those edits.
            fn for_module(
                root: &Path,
                module_name: &'a str,
                classes: &'a IndexMap<String, edits::Class>,
            ) -> anyhow::Result<Self> {
                let (input_path, input) = open_python_module(root, module_name)?;

                let tempfile_path = input_path.with_added_extension("tmp");
                let tempfile = BufWriter::new(File::create(&tempfile_path)?);

                Ok(Self {
                    context: PyiContext {
                        module_name,
                        classes,
                    },
                    state: PyiEditorState {
                        wrote_type_vars: false,
                        unseen_classes: classes.keys().map(String::as_str).collect(),
                        current_class: None,
                    },
                    input: FileInput { input_path, input },
                    output: FileOutput {
                        output_path: tempfile_path,
                        output: tempfile,
                    },
                })
            }
        }

        impl<'a> PyiEditorState<'a> {
            /// Once we've gotten to the end of a class, make sure that we've seen all the methods
            /// we're expecting and then clear the class state.
            fn finish_class(&mut self, context: &PyiContext<'a>) -> anyhow::Result<()> {
                let module_name = context.module_name;

                match self.current_class.take() {
                    Some(PyiCurrentClass {
                        class_name,
                        unseen_methods,
                        class: _,
                    }) if !unseen_methods.is_empty() => {
                        bail!(
                            "no type stubs found for methods: {}",
                            unseen_methods
                                .into_iter()
                                .map(|method| format!("{module_name}.{class_name}.{method}"))
                                .collect::<Vec<_>>()
                                .join(", ")
                        )
                    }

                    Some(PyiCurrentClass { .. }) | None => Ok(()),
                }
            }

            /// Given a line, print the resulting possibly-edited line and update this state.
            fn update_and_output<W: io::Write>(
                &mut self,
                context: &PyiContext<'a>,
                mut output: W,
                line: &str,
            ) -> anyhow::Result<()> {
                let PyiContext {
                    module_name,
                    classes,
                } = context;

                if line.chars().next().is_some_and(|c| !c.is_whitespace()) {
                    self.finish_class(context)?;

                    // We don't need to handle this case at the end of a module – if we found no
                    // classes, then we had no type variables to write out.
                    #[expect(
                        clippy::nonminimal_bool,
                        reason = "grouping the line prefixes is clearer"
                    )]
                    if !self.wrote_type_vars
                        && !(line.starts_with("import ")
                            || line.starts_with("from ")
                            || line.starts_with("#"))
                    {
                        write!(
                            output,
                            "{}",
                            edits::fmt_type_var_definitions(&edits::collect_all_type_parameters(
                                module_name,
                                classes
                            )?)
                        )?;

                        self.wrote_type_vars = true;
                    }
                }

                let Self {
                    wrote_type_vars: _,
                    unseen_classes,
                    current_class,
                } = self;

                if let Some((class_name, class)) = line
                    .strip_prefix("class ")
                    .and_then(|classless| classless.strip_suffix(":"))
                    .and_then(|class_name| classes.get_key_value(class_name))
                {
                    if !unseen_classes.shift_remove(class_name.as_str()) {
                        bail!("duplicate occurrences of class {module_name}.{class_name}");
                    }

                    *current_class = Some(PyiCurrentClass {
                        class_name,
                        class,
                        unseen_methods: class.methods.keys().map(String::as_str).collect(),
                    });

                    writeln!(
                        output,
                        "class {class_name}{}:",
                        edits::fmt_class_type_parameters_as_inheritance(&class.type_parameters)
                    )?;
                } else if let Some(PyiCurrentClass {
                    class_name,
                    class,
                    unseen_methods,
                }) = current_class.as_mut()
                {
                    let writeln_replacing_self = |output: &mut W, text: &str| {
                        let mut first = true;
                        for fragment in text.split("$SELF") {
                            if first {
                                first = false;
                            } else {
                                write!(
                                    output,
                                    "{class_name}{}",
                                    edits::fmt_type_parameters_defaultless(&class.type_parameters)
                                )?;
                            }
                            write!(output, "{fragment}")?;
                        }
                        writeln!(output)
                    };

                    if let Some((method_name, method, sig_no_lparen)) = line
                        .strip_prefix("    def ")
                        .and_then(|defless| defless.split_once("("))
                        .and_then(|(method_name, sig_no_lparen)| {
                            let method = class.methods.get(method_name)?;
                            Some((method_name, method, sig_no_lparen))
                        })
                    {
                        let edits::Method {
                            // We use old-style Python with explicit type vars, so we don't get to
                            // parameterize methods
                            type_parameters: _,
                            self_type,
                        } = method;

                        write!(output, "    def {method_name}(")?;

                        // Duplicate methods are fine – that's what `@overload` is
                        unseen_methods.shift_remove(method_name);

                        match self_type {
                            Some(self_type) => match sig_no_lparen.strip_prefix("self,") {
                                Some(selfless) => {
                                    write!(output, "self: {self_type},")?;
                                    writeln_replacing_self(&mut output, selfless)?;
                                }
                                None => bail!(
                                    "no self parameter for method \
                                     {module_name}.{class_name}.{method_name}"
                                ),
                            },
                            None => writeln_replacing_self(&mut output, sig_no_lparen)?,
                        }
                    } else {
                        writeln_replacing_self(&mut output, line)?;
                    }
                } else {
                    writeln!(output, "{line}")?;
                }

                Ok(())
            }
        }

        /// Given a path to the `root` where the Python `.pyi` files are stored, a Python `module`
        /// name, and a map of edits to be made to `classes` in that module, update the `.pyi` file
        /// for the given module.
        pub fn edit_module<'a>(
            root: &Path,
            module_name: &'a str,
            classes: &'a IndexMap<String, edits::Class>,
        ) -> anyhow::Result<()> {
            let PyiEditor {
                context,
                mut state,
                input: FileInput { input_path, input },
                output:
                    FileOutput {
                        output_path,
                        mut output,
                    },
            } = PyiEditor::for_module(root, module_name, classes)?;

            for line in input.lines() {
                state.update_and_output(&context, &mut output, &line?)?;
            }

            output.flush()?;
            drop(output);

            state.finish_class(&context)?;

            if !state.unseen_classes.is_empty() {
                let module_name = context.module_name;
                bail!(
                    "no type stubs found for classes: {}",
                    state
                        .unseen_classes
                        .into_iter()
                        .map(|class_name| format!("{module_name}.{class_name}"))
                        .collect::<Vec<_>>()
                        .join(", ")
                )
            }

            std::fs::rename(output_path, input_path)?;

            Ok(())
        }
    }

    /// The classes that need generic parameters added
    fn pyi_edits() -> edits::Modules {
        use edits::{Class, Method, Modules, TypeParameter};

        let real = TypeParameter::new("_Real");
        let complex = TypeParameter::defaulted("_Complex", "_Real");
        let ignored = TypeParameter::new("__T");

        let evaluable_with_extras =
            |class_name: &'static str, mut methods: IndexMap<String, Method>| {
                let None = methods.insert(
                    "evaluate".to_owned(),
                    Method {
                        type_parameters: vec![
                            TypeParameter::new("_OtherReal"),
                            TypeParameter::defaulted("_OtherComplex", "_OtherReal"),
                        ],
                        self_type: None,
                    },
                ) else {
                    panic!("tried to request multiple edits for {class_name}.evaluate");
                };

                (
                    class_name.to_owned(),
                    Class {
                        type_parameters: vec![real.clone(), complex.clone()],
                        methods,
                    },
                )
            };

        let builtin_waveform = |name| {
            evaluable_with_extras(
                name,
                IndexMap::from([(
                    "iq_values_at_sample_rate".to_owned(),
                    Method {
                        type_parameters: vec![ignored.clone()],
                        self_type: Some(format!("{name}[builtins.float, builtins.complex]")),
                    },
                )]),
            )
        };

        Modules::from([(
            "quil._quil.waveform".to_owned(),
            IndexMap::from([
                evaluable_with_extras(
                    "CommonBuiltinParameters",
                    IndexMap::from([(
                        "resolve_with_sample_rate".to_owned(),
                        Method {
                            type_parameters: vec![ignored.clone()],
                            self_type: Some(format!(
                                "CommonBuiltinParameters[builtins.float, {}]",
                                ignored.name
                            )),
                        },
                    )]),
                ),
                evaluable_with_extras("Waveform", IndexMap::new()),
                builtin_waveform("BuiltinWaveform"),
                builtin_waveform("Flat"),
                builtin_waveform("Gaussian"),
                builtin_waveform("DragGaussian"),
                builtin_waveform("ErfSquare"),
                builtin_waveform("HermiteGaussian"),
                (
                    "BoxcarKernel".to_owned(),
                    Class {
                        type_parameters: vec![],
                        methods: IndexMap::from([(
                            "iq_values_at_sample_rate".to_owned(),
                            Method {
                                type_parameters: vec![ignored.clone()],
                                self_type: None,
                            },
                        )]),
                    },
                ),
            ]),
        )])
    }

    pub fn main() -> anyhow::Result<()> {
        let mut stub = quil_rs::quilpy::stub_info()?;
        rigetti_pyo3::stubs::sort(&mut stub);
        stub.generate()?;
        for (module, classes) in pyi_edits() {
            editor::edit_module(&stub.python_root, &module, &classes)?;
        }
        Ok(())
    }
}

pub use main::main;