pyo3_stub_gen/generate/
variant_methods.rs

1use indexmap::IndexMap;
2
3use crate::generate::{MethodDef, MethodType, Parameter, ParameterDefault, Parameters};
4use crate::type_info::{ParameterKind, 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            parameters: Parameters::from_infos(info.constr_args),
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            is_overload: false,
32        });
33
34    if let VariantForm::Tuple = info.form {
35        let len_name = "__len__";
36        methods
37            .entry(len_name.to_string())
38            .or_default()
39            .push(MethodDef {
40                name: len_name,
41                parameters: Parameters::new(),
42                r#return: TypeInfo::builtin("int"),
43                doc: "",
44                r#type: MethodType::Instance,
45                is_async: false,
46                deprecated: None,
47                type_ignored: None,
48                is_overload: false,
49            });
50
51        let getitem_name = "__getitem__";
52        methods
53            .entry(getitem_name.to_string())
54            .or_default()
55            .push(MethodDef {
56                name: getitem_name,
57                parameters: Parameters {
58                    positional_or_keyword: vec![Parameter {
59                        name: "key",
60                        kind: ParameterKind::PositionalOrKeyword,
61                        type_info: TypeInfo::builtin("int"),
62                        default: ParameterDefault::None,
63                    }],
64                    ..Parameters::new()
65                },
66                r#return: TypeInfo::any(),
67                doc: "",
68                r#type: MethodType::Instance,
69                is_async: false,
70                deprecated: None,
71                type_ignored: None,
72                is_overload: false,
73            });
74    }
75
76    methods
77}
pyo3_stub_gen/generate/variant_methods.rs

pyo3_stub_gen/generate/
variant_methods.rs
