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            type_ignored: None,
31        });
32
33    if let VariantForm::Tuple = info.form {
34        let len_name = "__len__";
35        methods
36            .entry(len_name.to_string())
37            .or_default()
38            .push(MethodDef {
39                name: len_name,
40                args: Vec::new(),
41                r#return: TypeInfo::builtin("int"),
42                doc: "",
43                r#type: MethodType::Instance,
44                is_async: false,
45                deprecated: None,
46                type_ignored: None,
47            });
48
49        let getitem_name = "__getitem__";
50        methods
51            .entry(getitem_name.to_string())
52            .or_default()
53            .push(MethodDef {
54                name: getitem_name,
55                args: vec![Arg {
56                    name: "key",
57                    r#type: TypeInfo::builtin("int"),
58                    signature: None,
59                }],
60                r#return: TypeInfo::any(),
61                doc: "",
62                r#type: MethodType::Instance,
63                is_async: false,
64                deprecated: None,
65                type_ignored: None,
66            });
67    }
68
69    methods
70}
pyo3_stub_gen/generate/variant_methods.rs

pyo3_stub_gen/generate/
variant_methods.rs
