pyo3_stub_gen/generate/
variant_methods.rs

1use indexmap::IndexMap;
2
3use crate::generate::{Arg, MethodDef, MethodType};
4use crate::type_info::{PyComplexEnumInfo, VariantForm, VariantInfo};
5use crate::TypeInfo;
6use std::collections::HashSet;
7
8pub(super) fn get_variant_methods(
9    enum_info: &PyComplexEnumInfo,
10    info: &VariantInfo,
11) -> IndexMap<String, Vec<MethodDef>> {
12    let full_class_name = format!("{}.{}", enum_info.pyclass_name, info.pyclass_name);
13
14    let mut methods: IndexMap<String, Vec<MethodDef>> = IndexMap::new();
15
16    methods
17        .entry("__new__".to_string())
18        .or_default()
19        .push(MethodDef {
20            name: "__new__",
21            args: info.constr_args.iter().map(|a| a.into()).collect(),
22            r#return: TypeInfo {
23                name: full_class_name,
24                import: HashSet::new(),
25            },
26            doc: "",
27            r#type: MethodType::New,
28            is_async: false,
29            deprecated: None,
30        });
31
32    if let VariantForm::Tuple = info.form {
33        let len_name = "__len__";
34        methods
35            .entry(len_name.to_string())
36            .or_default()
37            .push(MethodDef {
38                name: len_name,
39                args: Vec::new(),
40                r#return: TypeInfo::builtin("int"),
41                doc: "",
42                r#type: MethodType::Instance,
43                is_async: false,
44                deprecated: None,
45            });
46
47        let getitem_name = "__getitem__";
48        methods
49            .entry(getitem_name.to_string())
50            .or_default()
51            .push(MethodDef {
52                name: getitem_name,
53                args: vec![Arg {
54                    name: "key",
55                    r#type: TypeInfo::builtin("int"),
56                    signature: None,
57                }],
58                r#return: TypeInfo::any(),
59                doc: "",
60                r#type: MethodType::Instance,
61                is_async: false,
62                deprecated: None,
63            });
64    }
65
66    methods
67}
pyo3_stub_gen/generate/variant_methods.rs

pyo3_stub_gen/generate/
variant_methods.rs
