pytorch/torch/csrc/jit/source_range.cpp
James Reed 6875018793 Record source/line info in SourceRange and report in highlight (#20898)
Summary:
Resolves https://github.com/pytorch/lockdown/issues/29

Examples:

```
import torch

torch.jit.script
def foobar(x):
    return torch.blargh(xyz)

==

RuntimeError:
object has no attribute blargh:
at compile.py:5:12
torch.jit.script
def foo(x):
    return torch.blargh(x)
           ~~~~~~~~~~~~ <--- HERE
```

It also gets the correct column number in the case where the original source file has common leading whitespace in front of the callable:

```
import torch

with torch.no_grad():
            torch.jit.script
            def foo(x):
                return torch.blargh(x)

==
RuntimeError:
object has no attribute blargh:
at compile_leading.py:6:24
torch.jit.script
def foo(x):
    return torch.blargh(x)
           ~~~~~~~~~~~~ <--- HERE
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20898

Differential Revision: D15552424

Pulled By: jamesr66a

fbshipit-source-id: 78d0f0de03f7ccbf3e7ea193a1b4eced57ea5d69
2019-05-29 21:32:33 -07:00

65 lines
2.2 KiB
C++

#include <torch/csrc/jit/source_range.h>
namespace torch {
namespace jit {
// a range of a shared string 'file_' with
C10_EXPORT void SourceRange::highlight(std::ostream& out) const {
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 << str;
return;
}
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')
--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')
++i;
if (i >= CONTEXT)
break;
}
AT_ASSERT(begin_highlight == 0 || str[begin_highlight - 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')
++i;
if (i >= CONTEXT)
break;
}
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());
out << std::string(len, '~')
<< (len < size() ? "... <--- HERE" : " <--- HERE");
out << str.substr(end_line, end_highlight - end_line);
if (!str.empty() && str.back() != '\n')
out << "\n";
}
} // namespace jit
} // namespace torch