Improve PythonException stack trace parser (#8002)

* Improve PythonException stack trace parser * Minor fix * Minor improvements
QuantConnect · May 9, 2024 · 28baf03 · 28baf03
1 parent 356b51f
commit 28baf03
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 41 deletions.
diff --git a/Common/Util/PythonUtil.cs b/Common/Util/PythonUtil.cs
@@ -21,6 +21,8 @@
 using QuantConnect.Data.Fundamental;
 using System.Text.RegularExpressions;
 using QuantConnect.Data.UniverseSelection;
+using System.IO;
+using System.Globalization;
 
 namespace QuantConnect.Util
 {
@@ -30,6 +32,7 @@ namespace QuantConnect.Util
     public class PythonUtil
     {
         private static Regex LineRegex = new Regex("line (\\d+)", RegexOptions.Compiled);
+        private static Regex StackTraceFileLineRegex = new Regex("\"(.+)\", line (\\d+), in (.+)", RegexOptions.Compiled | RegexOptions.Singleline);
         private static readonly Lazy<dynamic> lazyInspect = new Lazy<dynamic>(() => Py.Import("inspect"));
 
         /// <summary>
@@ -203,41 +206,50 @@ public static string PythonExceptionStackParser(string value)
                 return string.Empty;
             }
 
-            // Format the information in every line
-            var lines = value.Substring(1, value.Length - 1)
-                .Split(new[] { "  File " }, StringSplitOptions.RemoveEmptyEntries)
-                .Where(x => x.Split(',').Length > 2)
-                .Select(x =>
+            // The stack trace info before "at Python.Runtime." is the trace we want,
+            // which is for user Python code.
+            var endIndex = value.IndexOf("at Python.Runtime.", StringComparison.InvariantCulture);
+            var neededStackTrace = endIndex > 0 ? value.Substring(0, endIndex) : value;
+
+            // The stack trace is separated in blocks by file
+            var blocks = neededStackTrace.Split("  File ", StringSplitOptions.RemoveEmptyEntries)
+                .Select(fileTrace =>
                 {
-                    // Get the directory where the user files are located
-                    var baseScript = value.GetStringBetweenChars('\"', '\"');
-                    var length = Math.Max(baseScript.LastIndexOf('/'), baseScript.LastIndexOf('\\'));
-                    if (length < 0)
+                    var trimedTrace = fileTrace.Trim();
+                    if (string.IsNullOrWhiteSpace(trimedTrace))
+                    {
+                        return string.Empty;
+                    }
+
+                    var match = StackTraceFileLineRegex.Match(trimedTrace);
+                    if (!match.Success)
                     {
                         return string.Empty;
                     }
-                    var directory = baseScript.Substring(0, 1 + length);
 
-                    var info = x.Replace(directory, string.Empty).Split(',');
-                    var line = info[0].GetStringBetweenChars('\"', '\"');
-                    var lineNumber = int.Parse(info[1].Replace("line", string.Empty).Trim()) + ExceptionLineShift;
-                    line = $" in {line}: line {lineNumber}";
+                    var capture = match.Captures[0] as Match;
 
-                    info = info[2].Split(new[] { "\\n" }, StringSplitOptions.RemoveEmptyEntries);
-                    line = $" {info[0].Replace(" in ", " at ")}{line}";
+                    var filePath = capture.Groups[1].Value;
+                    var lastFileSeparatorIndex = Math.Max(filePath.LastIndexOf('/'), filePath.LastIndexOf('\\'));
+                    if (lastFileSeparatorIndex < 0)
+                    {
+                        return string.Empty;
+                    }
 
-                    // If we have the exact statement, add it to the error line
-                    if (info.Length > 2) line += $" :: {info[1].Trim()}";
+                    var fileName = filePath.Substring(lastFileSeparatorIndex + 1);
+                    var lineNumber = int.Parse(capture.Groups[2].Value, CultureInfo.InvariantCulture) + ExceptionLineShift;
+                    var locationAndInfo = capture.Groups[3].Value.Trim();
 
-                    return line;
-                });
+                    return $"  at {locationAndInfo}{Environment.NewLine} in {fileName}: line {lineNumber}";
+                })
+                .Where(x => !string.IsNullOrWhiteSpace(x));
 
-            var errorLine = string.Join(Environment.NewLine, lines);
-            errorLine = Extensions.ClearLeanPaths(errorLine);
+            var result = string.Join(Environment.NewLine, blocks);
+            result = Extensions.ClearLeanPaths(result);
 
-            return string.IsNullOrWhiteSpace(errorLine)
+            return string.IsNullOrWhiteSpace(result)
                 ? string.Empty
-                : $"{Environment.NewLine}{errorLine}{Environment.NewLine}";
+                : $"{Environment.NewLine}{result}{Environment.NewLine}";
         }
 
         /// <summary>

diff --git a/Tests/Common/Util/PythonUtilTests.cs b/Tests/Common/Util/PythonUtilTests.cs
@@ -121,9 +121,9 @@ public void ParsesPythonExceptionMessage(string expected, string original, int s
   at Initialize
     s
 ===
-   at Python.Runtime.PyObject.Invoke(PyTuple args in BasicTemplateAlgorithm.py: line 20
+ in BasicTemplateAlgorithm.py: line 20
 ",
-            @"   File ""D:/QuantConnect/MyLean/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 30, in Initialize
+            @"  File ""D:/QuantConnect/MyLean/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 30, in Initialize
     s
 ===
    at Python.Runtime.PyObject.Invoke(PyTuple args, PyDict kw)
@@ -136,12 +136,12 @@ at CallSite.Target(Closure , CallSite , Object )
             -10)]
         [TestCase(@"
   at Initialize
-    self.SetEndDate(201)
+    self.SetEndDate(201, 1)
 ===
-   at Python.Runtime.PyObject.Invoke(PyTuple args in BasicTemplateAlgorithm.py: line 40
+ in BasicTemplateAlgorithm.py: line 40
 ",
             @"  File ""D:/QuantConnect/MyLean/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 30, in Initialize
-    self.SetEndDate(201)
+    self.SetEndDate(201, 1)
 ===
    at Python.Runtime.PyObject.Invoke(PyTuple args, PyDict kw)
    at Python.Runtime.PyObject.InvokeMethod(String name, PyTuple args, PyDict kw)
@@ -154,11 +154,7 @@ at CallSite.Target(Closure , CallSite , Object )
         [TestCase(@"
   at <module>
     class BasicTemplateAlgorithm(QCAlgorithm):
-   at Python.Runtime.PythonException.ThrowLastAsClrException()
-   at Python.Runtime.NewReferenceExtensions.BorrowOrThrow(NewReference& reference)
-   at Python.Runtime.PyModule.Import(String name)
-   at Python.Runtime.Py.Import(String name)
-   at QuantConnect.AlgorithmFactory.Python.Wrappers.AlgorithmPythonWrapper..ctor(String moduleName) at AlgorithmFactory\Python\Wrappers\AlgorithmPythonWrapper.cs:line 74 in BasicTemplateAlgorithm.py: line 23
+ in BasicTemplateAlgorithm.py: line 23
 ",
             @"  File ""D:/QuantConnect/MyLean/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 23, in <module>
     class BasicTemplateAlgorithm(QCAlgorithm):
@@ -174,11 +170,10 @@ raise KeyError(f""No key found for either mapped or original key. Mapped Key: {m
  in PandasMapper.py: line 76
   at HistoryCall
     history.loc['QQQ']   # <--- raises Key Error
- in Algorithm.Python/TestAlgorithm.py: line 27
+ in TestAlgorithm.py: line 27
   at OnData
     self.HistoryCall()
-   at Python.Runtime.PythonException.ThrowLastAsClrException()
-   at Python.Runtime.PyObject.Invoke(PyTuple args in Algorithm.Python/TestAlgorithm.py: line 23
+ in TestAlgorithm.py: line 23
 ",
             @"  File ""/home/user/QuantConnect/Lean/Launcher/bin/Debug/PandasMapper.py"", line 76, in wrapped_function
     raise KeyError(f""No key found for either mapped or original key. Mapped Key: {mKey}; Original Key: {oKey}"")
@@ -191,14 +186,13 @@ at Python.Runtime.PyObject.Invoke(PyTuple args, PyDict kw)
    at Python.Runtime.PyObject.TryInvoke(InvokeBinder binder, Object[] args, Object& result)
    at CallSite.Target(Closure , CallSite , Object , PythonSlice )
    at System.Dynamic.UpdateDelegates.UpdateAndExecuteVoid2[T0,T1](CallSite site, T0 arg0, T1 arg1)
-   at QuantConnect.AlgorithmFactory.Python.Wrappers.AlgorithmPythonWrapper.OnData(Slice slice) in /home/jhonathan/QuantConnect/Lean/AlgorithmFactory/Python/Wrappers/AlgorithmPythonWrapper.cs:line 587
-   at QuantConnect.Lean.Engine.AlgorithmManager.Run(AlgorithmNodePacket job, IAlgorithm algorithm, ISynchronizer synchronizer, ITransactionHandler transactions, IResultHandler results, IRealTimeHandler realtime, ILeanManager leanManager, IAlphaHandler alphas, CancellationToken token) in /home/jhonathan/QuantConnect/Lean/Engine/AlgorithmManager.cs:line 523",
+   at QuantConnect.AlgorithmFactory.Python.Wrappers.AlgorithmPythonWrapper.OnData(Slice slice) in /home/user/QuantConnect/Lean/AlgorithmFactory/Python/Wrappers/AlgorithmPythonWrapper.cs:line 587
+   at QuantConnect.Lean.Engine.AlgorithmManager.Run(AlgorithmNodePacket job, IAlgorithm algorithm, ISynchronizer synchronizer, ITransactionHandler transactions, IResultHandler results, IRealTimeHandler realtime, ILeanManager leanManager, IAlphaHandler alphas, CancellationToken token) in /home/user/QuantConnect/Lean/Engine/AlgorithmManager.cs:line 523",
             0)]
         [TestCase(@"
   at OnData
     raise ValueError(""""ASD"""")
-   at Python.Runtime.PythonException.ThrowLastAsClrException() at src\runtime\PythonException.cs:line 52
-   at Python.Runtime.PyObject.Invoke(PyTuple args in BasicTemplateAlgorithm.py: line 43
+ in BasicTemplateAlgorithm.py: line 43
 ",
             @"  File ""D:\QuantConnect/MyLean/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 43, in OnData
     raise ValueError(""""ASD"""")
@@ -210,6 +204,23 @@ at CallSite.Target(Closure , CallSite , Object , PythonSlice )
    at QuantConnect.AlgorithmFactory.Python.Wrappers.AlgorithmPythonWrapper.OnData(Slice slice) in D:\QuantConnect\MyLean\Lean\AlgorithmFactory\Python\Wrappers\AlgorithmPythonWrapper.cs:line 693
    at QuantConnect.Lean.Engine.AlgorithmManager.Run(AlgorithmNodePacket job, IAlgorithm algorithm, ISynchronizer synchronizer, ITransactionHandler transactions, IResultHandler results, IRealTimeHandler realtime, ILeanManager leanManager, CancellationToken token) in D:\QuantConnect\MyLean\Lean\Engine\AlgorithmManager.cs:line 526",
             0)]
+        [TestCase(@"
+  at on_data
+    self.set_holdings(""SPY"", 1 / None)
+                             ~~^~~~~~
+ in BasicTemplateAlgorithm.py: line 46
+",
+            @" File ""C:\Users/user/QuantConnect/Lean/Algorithm.Python\BasicTemplateAlgorithm.py"", line 46, in on_data
+    self.set_holdings(""SPY"", 1 / None)
+                             ~~^~~~~~
+   at Python.Runtime.PythonException.ThrowLastAsClrException()
+   at Python.Runtime.PyObject.Invoke(PyTuple args, PyDict kw)
+   at Python.Runtime.PyObject.TryInvoke(InvokeBinder binder, Object[] args, Object& result)
+   at CallSite.Target(Closure , CallSite , Object , Slice )
+   at System.Dynamic.UpdateDelegates.UpdateAndExecuteVoid2[T0,T1](CallSite site, T0 arg0, T1 arg1)
+   at QuantConnect.AlgorithmFactory.Python.Wrappers.AlgorithmPythonWrapper.OnData(Slice slice) in C:\Users\user\QuantConnect\Lean\AlgorithmFactory\Python\Wrappers\AlgorithmPythonWrapper.cs:line 759
+   at QuantConnect.Lean.Engine.AlgorithmManager.Run(AlgorithmNodePacket job, IAlgorithm algorithm, ISynchronizer synchronizer, ITransactionHandler transactions, IResultHandler results, IRealTimeHandler realtime, ILeanManager leanManager, CancellationToken token) in C:\Users\user\QuantConnect\Lean\Engine\AlgorithmManager.cs:line 524",
+            0)]
         public void ParsesPythonExceptionStackTrace(string expected, string original, int shift)
         {
             var originalShiftValue = PythonUtil.ExceptionLineShift;