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

pyo3_stub_gen/generate/
variant_methods.rs
