From e84d7abcba6c93b348aa70540d67f79c47ea9b17 Mon Sep 17 00:00:00 2001
From: lojack5 <1458329+lojack5@users.noreply.github.com>
Date: Tue, 17 Oct 2023 21:34:31 -0600
Subject: [PATCH] Possible way to handle enums.

Named enums with 'Flags' in the name use `enum.IntFlag`, all other enums
use `enum.IntEnum`.  We have to bring the enum members into the containing
scope to match the original behaviour. Further, since these enums are typing
information only, and not actually in wx proper, we prevent them from appearing
to leak the the user by marking as non-exported (prepend '_' to the name), then
make a TypeAlias to the enum or an int.  This way type signatures still claim
to accept integers as appropriate.

I like this solution best, because we preserved the information about which
enum members are expected for method/function parameters, while not leaking
non-existant classes out of the type-stubs, and not complaining about using
integers.

There's still a few undefined but referenced enums (ex:
richtext.TextAttrDimensionFlags). These are most likely a union of some of
the other flags/enum types already defined, but the work to pull that information
from the C++ source is probably too much.
---
 etgtools/pi_generator.py  | 36 ++++++++++++++++++++++++++++++++++--
 etgtools/tweaker_tools.py |  1 +
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/etgtools/pi_generator.py b/etgtools/pi_generator.py
index 4ca49cec..1b3764a4 100644
--- a/etgtools/pi_generator.py
+++ b/etgtools/pi_generator.py
@@ -79,7 +79,8 @@ header_pyi = """\
 
 typing_imports = """\
 from __future__ import annotations
-from typing import Any, overload
+from enum import IntEnum, IntFlag, auto
+from typing import Any, overload, TypeAlias
 
 """
 
@@ -224,11 +225,42 @@ class PiWrapperGenerator(generators.WrapperGeneratorBase, FixWxPrefix):
         assert isinstance(enum, extractors.EnumDef)
         if enum.ignored or piIgnored(enum):
             return
+        # These enum classes aren't actually accessible from the real wx
+        # module, o we need to prepend _. But we want to make a type alias to
+        # the non-prefixed name, so method signatures can reference it without
+        # any special code, and also allow bare ints as inputs.
+        if '@' in enum.name or not enum.name:
+            # Anonymous enum
+            enum_name = f"_enum_{enum.name.replace('@', '').strip()}"
+            alias = ''
+        else:
+            alias = self.fixWxPrefix(enum.name)
+            enum_name = f'_{alias}'
+        if 'Flags' in enum_name:
+            enum_type = 'IntFlag'
+        else:
+            enum_type = 'IntEnum'
+        # Create the enum definition
+        stream.write(f'\n{indent}class {enum_name}({enum_type}):\n')
         for v in enum.items:
             if v.ignored or piIgnored(v):
                 continue
             name = v.pyName or v.name
-            stream.write('%s%s = 0\n' % (indent, name))
+            stream.write(f'{indent}    {name} = auto()\n')
+        # Create the alias if needed
+        if alias:
+            stream.write(f'{indent}{alias}: TypeAlias = {enum_name} | int\n')
+        # And bring the enum members into global scope.  We can't use
+        # enum.global_enum for this because:
+        #  1. It's only available on Python 3.11+
+        #  2. FixWxPrefix wouldn't be able to pick up the names, since it's
+        #     detecting based on AST parsing, not runtime changes (which
+        #     enum.global_enum performs).
+        for v in enum.items:
+            if v.ignored or piIgnored(v):
+                continue
+            name = v.pyName or v.name
+            stream.write(f'{indent}{name} = {enum_name}.{name}\n')
 
     #-----------------------------------------------------------------------
     def generateGlobalVar(self, globalVar, stream):
diff --git a/etgtools/tweaker_tools.py b/etgtools/tweaker_tools.py
index bce63049..3e6c1c8a 100644
--- a/etgtools/tweaker_tools.py
+++ b/etgtools/tweaker_tools.py
@@ -100,6 +100,7 @@ class FixWxPrefix(object):
         testName = name
         if '(' in name:
             testName = name[:name.find('(')]
+        testName = testName.split('.')[0]
 
         if testName in FixWxPrefix._coreTopLevelNames:
             return 'wx.'+name