pytorch · jamesr66a · May 30, 2019
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -3679,6 +3679,23 @@ def foo():
             return [[4]] + [[4, 5]]
         self.checkScript(foo, ())
 
+    def test_file_line_error(self):
+        def foobar(xyz):
+            return torch.blargh(xyz)
+
+        _, lineno = inspect.getsourcelines(foobar)
+        with self.assertRaisesRegex(RuntimeError, "test_jit.py:{}:20".format(lineno + 1)):
+            scripted = torch.jit.script(foobar)
+
+    def test_file_line_error_class_defn(self):
+        class FooBar(object):
+            def baz(self, xyz):
+                return torch.blargh(xyz)
+
+        _, lineno = inspect.getsourcelines(FooBar)
+        with self.assertRaisesRegex(RuntimeError, "test_jit.py:{}:24".format(lineno + 2)):
+            torch.jit.script(FooBar)
+
     def test_tensor_shape(self):
         x = torch.empty(34, 56, 78)
 

diff --git a/torch/csrc/jit/ir.cpp b/torch/csrc/jit/ir.cpp
@@ -1231,8 +1231,7 @@ void Node::removeFromList() {
 }
 
 inline const SourceRange& fakeRange() {
-  static SourceRange range(
-      std::make_shared<std::string>("<internally-created-node>"), 0, 1);
+  static SourceRange range(std::make_shared<Source>(""), 0, 1);
   return range;
 }
 

diff --git a/torch/csrc/jit/script/lexer.h b/torch/csrc/jit/script/lexer.h
@@ -368,7 +368,7 @@ struct Token {
 
 struct Lexer {
   explicit Lexer(const std::string& str)
-      : file(std::make_shared<std::string>(str)),
+      : source(std::make_shared<Source>(str)),
         pos(0),
         nesting(0),
         indent_stack(),
@@ -485,9 +485,9 @@ struct Lexer {
     int kind;
     size_t start;
     size_t length;
-    AT_ASSERT(file);
+    AT_ASSERT(source);
     if (!shared.match(
-            *file,
+            source->text(),
             pos,
             nesting > 0,
             whitespace_token,
@@ -496,14 +496,15 @@ struct Lexer {
             &length)) {
       expected(
           "a valid token",
-          Token((*file)[start], SourceRange(file, start, start + 1)));
+          Token(
+              (source->text())[start], SourceRange(source, start, start + 1)));
     }
-    auto t = Token(kind, SourceRange(file, start, start + length));
+    auto t = Token(kind, SourceRange(source, start, start + length));
     pos = start + length;
     return t;
   }
 
-  std::shared_ptr<std::string> file;
+  std::shared_ptr<Source> source;
   size_t pos;
   size_t nesting; // depth of ( [ { nesting...
   std::vector<int> indent_stack; // stack of identation level of blocks

diff --git a/torch/csrc/jit/script/python_tree_views.cpp b/torch/csrc/jit/script/python_tree_views.cpp
@@ -14,27 +14,44 @@ namespace jit {
 namespace script {
 
 struct SourceRangeFactory {
-  SourceRangeFactory(std::string source)
-      : source_(std::make_shared<std::string>(std::move(source))) {
-    size_t pos = 0;
-    do {
-      line_len_prefix_sum_.push_back(pos);
-      pos++;
-    } while ((pos = source_->find('\n', pos)) != std::string::npos);
-  }
+  SourceRangeFactory(
+      std::string text,
+      std::string filename,
+      size_t file_lineno,
+      size_t leading_whitespace_chars)
+      : source_(std::make_shared<Source>(
+            std::move(text),
+            std::move(filename),
+            file_lineno)),
+        leading_whitespace_chars_(leading_whitespace_chars) {}
   SourceRange create(int line, int start_col, int end_col) {
+    size_t start_byte_offset, end_byte_offset;
+    std::tie(start_byte_offset, end_byte_offset) =
+        line_col_to_byte_offs(
+            line,
+            start_col + leading_whitespace_chars_,
+            end_col + leading_whitespace_chars_);
+    return SourceRange(source_, start_byte_offset, end_byte_offset);
+  }
+
+  std::tuple<size_t, size_t> line_col_to_byte_offs(
+      int line,
+      int start_col,
+      int end_col) {
     // Python has a weird convention where col_offset points to the column
     // *before* the token starts.
     start_col++;
     end_col++;
     // Also, lines are counted from 1.
     line--;
-    auto line_start = line_len_prefix_sum_.at(line);
-    return SourceRange(source_, line_start + start_col, line_start + end_col);
+    auto line_start = source_->offset_for_line(line);
+    return std::make_tuple<size_t, size_t>(
+        line_start + start_col, line_start + end_col);
   }
 
-  std::shared_ptr<std::string> source_;
+  std::shared_ptr<Source> source_;
   std::vector<size_t> line_len_prefix_sum_;
+  size_t leading_whitespace_chars_;
 };
 
 template <typename T>
@@ -65,15 +82,15 @@ void initTreeViewBindings(PyObject* module) {
       .def_property_readonly("start", &SourceRange::start)
       .def_property_readonly("end", &SourceRange::end);
   py::class_<SourceRangeFactory>(m, "SourceRangeFactory")
-      .def(py::init<std::string&&>())
+      .def(py::init<std::string&&, std::string&&, size_t, size_t>())
       .def("make_range", &SourceRangeFactory::create)
       .def(
           "make_raw_range",
           [](const SourceRangeFactory& self, size_t start, size_t end) {
             return SourceRange(self.source_, start, end);
           })
       .def_property_readonly("source", [](const SourceRangeFactory& self) {
-        return *self.source_;
+        return self.source_->text();
       });
 
   py::class_<TreeView>(m, "TreeView")

diff --git a/torch/csrc/jit/script/tree.h b/torch/csrc/jit/script/tree.h
@@ -128,7 +128,7 @@ static SourceRange mergeRanges(SourceRange c, const TreeList& others) {
       continue;
     size_t s = std::min(c.start(), t->range().start());
     size_t e = std::max(c.end(), t->range().end());
-    c = SourceRange(c.file_ptr(), s, e);
+    c = SourceRange(c.source(), s, e);
   }
   return c;
 }

diff --git a/torch/csrc/jit/source_range.cpp b/torch/csrc/jit/source_range.cpp
@@ -5,14 +5,14 @@ namespace jit {
 
 // a range of a shared string 'file_' with
 C10_EXPORT void SourceRange::highlight(std::ostream& out) const {
-  if (size() == file_->size()) {
+  const std::string& str = source_->text();
+  if (size() == str.size()) {
     // this is just the entire file, not a subset, so print it out.
     // primarily used to print out python stack traces
-    out << *file_;
+    out << str;
     return;
   }
 
-  const std::string& str = file();
   size_t begin_line = start(); // beginning of line to highlight
   size_t end_line = start(); // end of line to highlight
   while (begin_line > 0 && str[begin_line - 1] != '\n')
@@ -42,6 +42,14 @@ C10_EXPORT void SourceRange::highlight(std::ostream& out) const {
   }
   AT_ASSERT(end_highlight == str.size() || str[end_highlight] == '\n');
 
+  if (source_->filename()) {
+    auto lineno = source_->lineno_for_offset(start());
+    auto col_offset = (int)start() -
+        (int)source_->offset_for_line(lineno);
+    out << "at " << *source_->filename() << ":"
+        << source_->lineno_to_source_lineno(lineno) << ":" << col_offset
+        << "\n";
+  }
   out << str.substr(begin_highlight, end_line - begin_highlight) << "\n";
   out << std::string(start() - begin_line, ' ');
   size_t len = std::min(size(), end_line - start());

diff --git a/torch/csrc/jit/source_range.h b/torch/csrc/jit/source_range.h
@@ -1,36 +1,109 @@
 #pragma once
 #include <c10/util/Exception.h>
+#include <c10/util/Optional.h>
 
 #include <algorithm>
 #include <memory>
 #include <iostream>
 namespace torch {
 namespace jit {
 
-// a range of a shared string 'file_' with functions to help debug by highlight
-// that
-// range.
+// Source represents a code segment. It keeps track of:
+//  - text : the text of the code segment
+//  - filename (optional) : if present, represents the name of the file from
+//                          which the code semgemnt originated.
+//  - starting_line_no : represents the line in the original file where the
+//                       code segment started.
+struct Source {
+  explicit Source(std::string text)
+      : text_(std::move(text)), filename_(c10::nullopt) {
+    calc_line_start_offsets();
+  }
+
+  Source(
+      std::string text,
+      c10::optional<std::string> filename,
+      size_t starting_line_no)
+      : text_(std::move(text)),
+        filename_(std::move(filename)),
+        starting_line_no_(starting_line_no) {
+    calc_line_start_offsets();
+  }
+
+  // Given a line number (within source_), return the byte offset of the
+  // beginning of that line.
+  size_t offset_for_line(size_t line) const {
+    return line_starting_offsets_.at(line);
+  }
+
+  // Calculate the line (within the code segment) on which `offset` resides.
+  size_t lineno_for_offset(size_t offset) const {
+    return std::upper_bound(
+               line_starting_offsets_.begin(),
+               line_starting_offsets_.end(),
+               offset) -
+        line_starting_offsets_.begin() - 1;
+  }
+
+  // Calculate the line (within the original source file, if present) on which
+  // `lineno` resides.
+  size_t lineno_to_source_lineno(size_t lineno) const {
+    if (filename_) {
+      return lineno + starting_line_no_;
+    } else {
+      return lineno;
+    }
+  }
+
+  const std::string& text() const {
+    return text_;
+  }
+
+  const c10::optional<std::string>& filename() const {
+    return filename_;
+  }
+
+  size_t starting_line_no() const {
+    return starting_line_no_;
+  }
+
+ private:
+  void calc_line_start_offsets() {
+    size_t pos = 0;
+    do {
+      line_starting_offsets_.push_back(pos);
+      pos++;
+    } while ((pos = text_.find('\n', pos)) != std::string::npos);
+  }
+  std::string text_;
+  c10::optional<std::string> filename_;
+  // If filename_ is not present, starting_line_no_ is don't care
+  size_t starting_line_no_;
+  // Starting offsets for lines into the source. e.g. line 0 starts at
+  // line_starting_offsets_[0], etc.
+  std::vector<size_t> line_starting_offsets_;
+};
+
+// A SourceRange is a view into a Source, that points to a subset of the source,
+// specified by `start` and `end` byte offsets into the source text.
 struct CAFFE2_API SourceRange {
-  SourceRange(std::shared_ptr<std::string> file_, size_t start_, size_t end_)
-      : file_(std::move(file_)), start_(start_), end_(end_) {}
+  SourceRange(std::shared_ptr<Source> source_, size_t start_, size_t end_)
+      : source_(std::move(source_)), start_(start_), end_(end_) {}
   explicit SourceRange(std::string string_range)
-      : file_(std::make_shared<std::string>(std::move(string_range))),
+      : source_(std::make_shared<Source>(std::move(string_range))),
         start_(0),
-        end_(file_->size()) {}
+        end_(source_->text().size()) {}
 
   const std::string text() const {
-    return file().substr(start(), end() - start());
+    return source_->text().substr(start(), end() - start());
   }
   size_t size() const {
     return end() - start();
   }
   static const size_t CONTEXT = 10;
   void highlight(std::ostream& out) const;
-  const std::string& file() const {
-    return *file_;
-  }
-  const std::shared_ptr<std::string>& file_ptr() const {
-    return file_;
+  const std::shared_ptr<Source>& source() const {
+    return source_;
   }
   size_t start() const {
     return start_;
@@ -68,7 +141,7 @@ struct CAFFE2_API SourceRange {
   }
 
  private:
-  std::shared_ptr<std::string> file_;
+  std::shared_ptr<Source> source_;
   size_t start_;
   size_t end_;
 };