From 62e21b308102f2966add6a5f35509c1c1c37fffe Mon Sep 17 00:00:00 2001
From: Martin Molinero <martin.molinero1@gmail.com>
Date: Thu, 5 Sep 2024 20:25:05 -0300
Subject: [PATCH] Use python Enum

- Minor adjustment so we use Python Enum class instead of CSharp system
  enum for stubs, it's better supported by mypy
---
 QuantConnectStubsGenerator/Renderer/ClassRenderer.cs     | 9 +++++++++
 QuantConnectStubsGenerator/Renderer/NamespaceRenderer.cs | 2 ++
 integration/integration_tests.py                         | 2 +-
 3 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/QuantConnectStubsGenerator/Renderer/ClassRenderer.cs b/QuantConnectStubsGenerator/Renderer/ClassRenderer.cs
index e189f13..08e0fae 100644
--- a/QuantConnectStubsGenerator/Renderer/ClassRenderer.cs
+++ b/QuantConnectStubsGenerator/Renderer/ClassRenderer.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -43,6 +44,14 @@ private void RenderClassHeader(Class cls)
 
             if (inherited.Count > 0)
             {
+                for (var i = 0; i < inherited.Count; i++)
+                {
+                    if (inherited[i].Equals("System.Enum", StringComparison.InvariantCultureIgnoreCase))
+                    {
+                        // 'Enum' is a python base type which is handled better by mypy if we used 'System' it assumes the enum value and causes a warning/missmatch
+                        inherited[i] = "Enum";
+                    }
+                }
                 Write($"({string.Join(", ", inherited)})");
             }
 
diff --git a/QuantConnectStubsGenerator/Renderer/NamespaceRenderer.cs b/QuantConnectStubsGenerator/Renderer/NamespaceRenderer.cs
index 0a9562d..bd54906 100644
--- a/QuantConnectStubsGenerator/Renderer/NamespaceRenderer.cs
+++ b/QuantConnectStubsGenerator/Renderer/NamespaceRenderer.cs
@@ -17,6 +17,8 @@ public override void Render(Namespace ns)
         {
             // Fix for Jedi; Include import of typing instead of using typing.overload
             WriteLine("from typing import overload");
+            // fix for python enums
+            WriteLine("from enum import Enum");
 
             var usedTypes = ns
                 .GetParentClasses()
diff --git a/integration/integration_tests.py b/integration/integration_tests.py
index f82e192..d1f908e 100644
--- a/integration/integration_tests.py
+++ b/integration/integration_tests.py
@@ -34,7 +34,7 @@ def main():
         file.write(f"""
 {{
     "include": [{", ".join([f'"{ns}/**"' for ns in os.listdir(stubs_dir)])}],
-    "exclude": ["System/Collections/Immutable/**", "System/__init__.pyi", "System/Runtime/InteropServices/__init__.pyi"],
+    "exclude": ["System/Collections/Immutable/**", "System/__init__.pyi", "System/Runtime/InteropServices/**"],
     "reportGeneralTypeIssues": "none",
     "reportInvalidTypeVarUse": "none",
     "reportWildcardImportFromLibrary": "none"