up

Author: blueloveTH

stubgen/enum.py

@@ -2,24 +2,25 @@
 from .writer import Writer
 from typing import Any
 
-def _gen_raw_enums(w: Writer, name: str, values: list[tuple[str, Any]], is_bitfield: bool) -> None:
-    if is_bitfield:
-        w.write(f"{name} = int")
-    else:
-        literal_values = [str(v) for _, v in values]
-        w.write(f"{name} = Literal[{', '.join(literal_values)}]")
-
-    for k, v in values:
-        w.write(f"{k} = {v}")
+def _gen_raw_enums(w: Writer, name: str, values: list[tuple[str, Any]], is_bitfield: bool, with_types=True, with_values=True) -> None:
+    if with_types:
+        if is_bitfield:
+            w.write(f"{name} = int")
+        else:
+            literal_values = [str(v) for _, v in values]
+            w.write(f"{name} = Literal[{', '.join(literal_values)}]")
+    if with_values:
+        for k, v in values:
+            w.write(f"{k} = {v}")
     w.write('')
 
 
-def gen_global_enums(w: Writer, global_enums: list[GlobalEnum]) -> None:
+def gen_global_enums(w: Writer, global_enums: list[GlobalEnum], with_types: bool, with_values: bool) -> None:
     for e in global_enums:
         name = e.name.replace('.', '_')
         _gen_raw_enums(w, name, [
             (v.name, v.value) for v in e.values
-        ], e.is_bitfield)
+        ], e.is_bitfield, with_types, with_values)
 
 def gen_class_enum(w: Writer, e: ClassesEnum):
     assert e.values is not None

stubgen/map.py

@@ -203,7 +203,6 @@ def gen_typings_pyi_writers(gdt_all_in_one: GodotInOne):
     writer_v.write(
         """\
 from typing import overload, Literal, Any
-from .enums import *
 from .classes._Object import Object
 from . import alias
 
@@ -246,6 +245,7 @@ def default(gdt_expr: str) -> Any: ...
             writer.write(
         """\
 from typing import Literal, Callable as typing_Callable
+from ..enums import *
 from ._init import *
 
 """
@@ -551,7 +551,6 @@ def gen_alias_pyi_writer(gdt_all_in_one: GodotInOne, pyi_writer: Writer) -> Writ
 def gen_init_pyi_writer(gdt_all_in_one: GodotInOne, pyi_writer: Writer, global_variant_classes: list[str]) -> Writer:
     pyi_writer.write(
         """\
-from .enums import *
 from .header import *
 from . import classes
 
@@ -567,6 +566,9 @@ def load(path: str) -> classes.Resource: ...
 
     for clazz in global_variant_classes:
         writer.writefmt('from .variants import {0} as {0}', clazz)
+    writer.write('')
+
+    enum.gen_global_enums(pyi_writer, gdt_all_in_one.global_enums, with_types=False, with_values=True)
     return pyi_writer
 
 
@@ -577,7 +579,7 @@ def gen_enums_pyi_writer(gdt_all_in_one: GodotInOne, pyi_writer: Writer) -> Writ
 """
     )
     # --- global enum ---
-    enum.gen_global_enums(pyi_writer, gdt_all_in_one.global_enums)
+    enum.gen_global_enums(pyi_writer, gdt_all_in_one.global_enums, with_types=True, with_values=False)
     return pyi_writer