Error printing extension support for multiline errors (#43807)

Summary: Pull Request resolved: pytorch/pytorch#43807

Test Plan: Imported from OSS

Reviewed By: gmagogsfm

Differential Revision: D23407457

Pulled By: Lilyjjo

fbshipit-source-id: 05a6a50dc39c00474d9087ef56028a2c183aa53a

docs/source/jit_language_reference.rst

@@ -102,16 +102,19 @@ Example (a type mismatch)
      @torch.jit.script
      def an_error(x):
          if x:
-         ~~~~~...  <--- HERE
+         ~~~~~
              r = torch.rand(1)
+             ~~~~~~~~~~~~~~~~~
          else:
+         ~~~~~
+             r = 4
+             ~~~~~ <--- HERE
+         return r
      and was used here:
          else:
              r = 4
          return r
-                ~ <--- HERE
-     ...
-
+                ~ <--- HERE...
 
 Unsupported Typing Constructs
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -738,10 +741,15 @@ Example:
      @torch.jit.script...
      def foo(x):
          if x < 0:
-         ~~~~~~~~~...  <--- HERE
+         ~~~~~~~~~
+             y = 4
+             ~~~~~ <--- HERE
+         print(y)
+     and was used here:
+         if x < 0:
              y = 4
          print(y)
-     ...
+               ~ <--- HERE...
 
 Non-local variables are resolved to Python values at compile time when the
 function is defined. These values are then converted into TorchScript values using

test/test_jit.py

@@ -5677,8 +5677,9 @@ def fn():
         ast = torch.jit.frontend.get_jit_def(fn, fn.__name__)
         FileCheck().check("SourceRange at:") \
                    .check("def fn():") \
-                   .check("~~~~~~~~~...  <--- HERE") \
+                   .check("~~~~~~~~~") \
                    .check('raise Exception("hello")') \
+                   .check('~~~~~~~~~~~~~~~~~ <--- HERE') \
                    .run(str(ast.range()))
 
     def test_python_frontend_py3(self):

torch/csrc/jit/frontend/source_range.cpp

@@ -75,37 +75,47 @@ C10_EXPORT void SourceRange::print_with_context(
     return;
   }
 
-  size_t begin_line = start(); // beginning of line to highlight
-  size_t end_line = start(); // end of line to highlight
+  size_t range_end =
+      (str.size() < end()
+           ? str.size()
+           : end()); // use instead of 'end()' because some ranges extend past
+                     // the length of the source
+
+  // determine CONTEXT line range
+  size_t begin_line = start(); // beginning of lines to highlight
+  size_t end_line = range_end;
   while (begin_line > 0 && str[begin_line - 1] != '\n')
     --begin_line;
   while (end_line < str.size() && str[end_line] != '\n')
     ++end_line;
   AT_ASSERT(begin_line == 0 || str[begin_line - 1] == '\n');
   AT_ASSERT(end_line == str.size() || str[end_line] == '\n');
 
-  size_t begin_highlight = begin_line; // beginning of context, CONTEXT lines
-                                       // before the highlight line
-  for (size_t i = 0; begin_highlight > 0; --begin_highlight) {
-    if (str[begin_highlight - 1] == '\n')
+  size_t begin_context = begin_line; // beginning of context, CONTEXT lines
+                                     // before the highlight lines
+  for (size_t i = 0; begin_context > 0; --begin_context) {
+    if (str[begin_context - 1] == '\n') {
       ++i;
+    }
     if (i >= context) {
       break;
     }
   }
-  AT_ASSERT(begin_highlight == 0 || str[begin_highlight - 1] == '\n');
+  AT_ASSERT(begin_context == 0 || str[begin_context - 1] == '\n');
 
-  size_t end_highlight =
-      end_line; // end of context, CONTEXT lines after the highlight line
-  for (size_t i = 0; end_highlight < str.size(); ++end_highlight) {
-    if (str[end_highlight] == '\n')
+  size_t end_context =
+      end_line; // end of context, CONTEXT lines after the highlight lines
+  for (size_t i = 0; end_context < str.size(); ++end_context) {
+    if (str[end_context] == '\n') {
       ++i;
+    }
     if (i >= context) {
       break;
     }
   }
-  AT_ASSERT(end_highlight == str.size() || str[end_highlight] == '\n');
+  AT_ASSERT(end_context == str.size() || str[end_context] == '\n');
 
+  // print out location information
   if (auto flc = file_line_col()) {
     std::string filename;
     size_t line, col;
@@ -116,17 +126,74 @@ C10_EXPORT void SourceRange::print_with_context(
     }
     out << "\n";
   }
-  out << str.substr(begin_highlight, end_line - begin_highlight) << "\n";
+  // print out inital context
+  out << str.substr(begin_context, start() - begin_context);
+  size_t line_start = start();
+  size_t line_end = range_end;
   if (highlight) {
-    out << std::string(start() - begin_line, ' ');
-    size_t len = std::min(size(), end_line - start());
-    out << std::string(len, '~')
-        << (len < size() ? "...  <--- HERE" : " <--- HERE");
+    line_end = start();
+    while (line_start < range_end) {
+      // move line_end to end of line
+      while (str[line_end] != '\n' && line_end < str.size()) {
+        ++line_end;
+      }
+      // print line of code
+      auto actual_line = str.substr(line_start, (line_end - line_start) + 1);
+      out << actual_line;
+      if (actual_line.back() != '\n') {
+        out << "\n";
+      }
+
+      size_t empty_space = 0;
+      size_t highlight_space = 0;
+      size_t hightlight_begin = line_start;
+      size_t highlight_end = line_start;
+      // determine length of line which is being highlighted
+      while (hightlight_begin > 0 && str[hightlight_begin - 1] != '\n') {
+        --hightlight_begin;
+      }
+      while (highlight_end < range_end && str[highlight_end] != '\n') {
+        ++highlight_end;
+      }
+      AT_ASSERT(hightlight_begin == 0 || str[hightlight_begin - 1] == '\n');
+      AT_ASSERT(highlight_end == range_end || str[highlight_end] == '\n');
+      // determine amount of empty space vs highlighted space
+      for (size_t i = hightlight_begin; i < highlight_end; i++) {
+        if (str[i] == ' ' || i < start()) {
+          empty_space++;
+        } else {
+          break;
+        }
+      }
+      highlight_space = highlight_end - hightlight_begin - empty_space;
+      if (highlight_space > 0) {
+        // some ranges are off and include empty white space on new lines which
+        // don't need to be printed
+        bool more_lines = false;
+        for (size_t i = line_end; i <= range_end; i++) {
+          if (str[i] != '\n' && str[i] != ' ') {
+            more_lines = true;
+          }
+        }
+        out << std::string(empty_space, ' ');
+        out << std::string(highlight_space, '~');
+        out << (more_lines && line_end != range_end ? "\n" : " <--- HERE\n");
+      }
+      ++line_end;
+      line_start = line_end;
+    }
+  } else {
+    // print out code with no highlight
+    out << str.substr(start(), range_end - start());
+  }
+  // print out ending context
+  if (line_end <= str.size()) {
+    auto line_substr = str.substr(line_end, end_context - line_end);
+    out << line_substr;
+    if (!line_substr.empty() && line_substr.back() != '\n') {
+      out << "\n";
+    }
   }
-  auto line_substr = str.substr(end_line, end_highlight - end_line);
-  out << line_substr;
-  if (!line_substr.empty() && line_substr.back() != '\n')
-    out << "\n";
 }
 
 } // namespace jit