You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2022/08/11 05:44:00 UTC
[tvm] branch main updated: [TVMScript] Text underlining in DocPrinter based on Doc's source_paths (#12344)
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f5f5a75ae9 [TVMScript] Text underlining in DocPrinter based on Doc's source_paths (#12344)
f5f5a75ae9 is described below
commit f5f5a75ae929535ccb59698c1ab83115d16d8bcc
Author: Greg Bonik <gb...@octoml.ai>
AuthorDate: Wed Aug 10 22:43:54 2022 -0700
[TVMScript] Text underlining in DocPrinter based on Doc's source_paths (#12344)
This adds an ability to print a "diagnostic marker" based on a given ObjectPath. For example, say we are printing a fragment of TIR like
```
for i in T.serial(10):
a[i] = 5
```
and we would like bring the user's attention to the bound of the loop:
```
for i in T.serial(10):
^^
a[i] = 5
```
In this case we would give the doc printer an object path that represents this loop bound, i.e. something like `path_to_underline=ObjectPath.root().attr("extent")`
Tracking issue: https://github.com/apache/tvm/issues/11912
---
include/tvm/script/printer/doc_printer.h | 11 +-
python/tvm/script/printer/doc_printer.py | 22 +-
src/script/printer/base_doc_printer.cc | 261 ++++++++++++++-
src/script/printer/base_doc_printer.h | 64 +++-
src/script/printer/python_doc_printer.cc | 16 +-
.../unittest/test_tvmscript_printer_underlining.py | 361 +++++++++++++++++++++
6 files changed, 718 insertions(+), 17 deletions(-)
diff --git a/include/tvm/script/printer/doc_printer.h b/include/tvm/script/printer/doc_printer.h
index 6bf502fab9..04a67a9b82 100644
--- a/include/tvm/script/printer/doc_printer.h
+++ b/include/tvm/script/printer/doc_printer.h
@@ -31,10 +31,15 @@ namespace printer {
* This function unpacks the DocPrinterOptions into function arguments
* to be FFI friendly.
*
- * \param doc the doc to be converted
- * \param indent_spaces the number of spaces used for indention
+ * \param doc Doc to be converted
+ * \param indent_spaces Number of spaces used for indentation
+ * \param print_line_numbers Whether to print line numbers
+ * \param num_context_lines Number of context lines to print around the underlined text
+ * \param path_to_underline Object path to be underlined
*/
-String DocToPythonScript(Doc doc, int indent_spaces = 4);
+String DocToPythonScript(Doc doc, int indent_spaces = 4, bool print_line_numbers = false,
+ int num_context_lines = -1,
+ Optional<ObjectPath> path_to_underline = NullOpt);
} // namespace printer
} // namespace script
diff --git a/python/tvm/script/printer/doc_printer.py b/python/tvm/script/printer/doc_printer.py
index 1cb56ecbf7..1791f46b00 100644
--- a/python/tvm/script/printer/doc_printer.py
+++ b/python/tvm/script/printer/doc_printer.py
@@ -16,11 +16,19 @@
# under the License.
"""Functions to print doc into text format"""
+from typing import Optional
+from tvm.runtime.object_path import ObjectPath
from . import _ffi_api
from .doc import Doc
-def to_python_script(doc: Doc, indent_spaces: int = 4) -> str:
+def to_python_script(
+ doc: Doc,
+ indent_spaces: int = 4,
+ print_line_numbers: bool = False,
+ num_context_lines: Optional[int] = None,
+ path_to_underline: Optional[ObjectPath] = None,
+) -> str:
"""Convert Doc into Python script.
Parameters
@@ -29,10 +37,20 @@ def to_python_script(doc: Doc, indent_spaces: int = 4) -> str:
The doc to convert into Python script
indent_spaces : int
The number of indent spaces to use in the output
+ print_line_numbers: bool
+ Whether to print line numbers
+ num_context_lines : Optional[int]
+ Number of context lines to print around the underlined text
+ path_to_underline : Optional[ObjectPath]
+ Object path to be underlined
Returns
-------
script : str
The text representation of Doc in Python syntax
"""
- return _ffi_api.DocToPythonScript(doc, indent_spaces) # type: ignore # pylint: disable=no-member
+ if num_context_lines is None:
+ num_context_lines = -1
+ return _ffi_api.DocToPythonScript( # type: ignore
+ doc, indent_spaces, print_line_numbers, num_context_lines, path_to_underline
+ )
diff --git a/src/script/printer/base_doc_printer.cc b/src/script/printer/base_doc_printer.cc
index 4129152129..38b8ef8977 100644
--- a/src/script/printer/base_doc_printer.cc
+++ b/src/script/printer/base_doc_printer.cc
@@ -23,19 +23,256 @@ namespace tvm {
namespace script {
namespace printer {
-DocPrinter::DocPrinter(int indent_spaces) : indent_spaces_(indent_spaces) {}
+namespace {
-void DocPrinter::Append(const Doc& doc) { PrintDoc(doc); }
+void SortAndMergeSpans(std::vector<ByteSpan>* spans) {
+ if (spans->empty()) {
+ return;
+ }
+ std::sort(spans->begin(), spans->end());
+ auto last = spans->begin();
+ for (auto cur = spans->begin() + 1; cur != spans->end(); ++cur) {
+ if (cur->first > last->second) {
+ *++last = *cur;
+ } else if (cur->second > last->second) {
+ last->second = cur->second;
+ }
+ }
+ spans->erase(++last, spans->end());
+}
+
+size_t GetTextWidth(const std::string& text, const ByteSpan& span) {
+ // FIXME: this only works for ASCII characters.
+ // To do this "correctly", we need to parse UTF-8 into codepoints
+ // and call wcwidth() or equivalent for every codepoint.
+ size_t ret = 0;
+ for (size_t i = span.first; i != span.second; ++i) {
+ if (isprint(text[i])) {
+ ret += 1;
+ }
+ }
+ return ret;
+}
+
+size_t MoveBack(size_t pos, size_t distance) { return distance > pos ? 0 : pos - distance; }
+
+size_t MoveForward(size_t pos, size_t distance, size_t max) {
+ return distance > max - pos ? max : pos + distance;
+}
+
+size_t GetLineIndex(size_t byte_pos, const std::vector<size_t>& line_starts) {
+ auto it = std::upper_bound(line_starts.begin(), line_starts.end(), byte_pos);
+ return (it - line_starts.begin()) - 1;
+}
+
+using UnderlineIter = typename std::vector<ByteSpan>::const_iterator;
+
+ByteSpan PopNextUnderline(UnderlineIter* next_underline, UnderlineIter end_underline) {
+ if (*next_underline == end_underline) {
+ return {std::numeric_limits<size_t>::max(), std::numeric_limits<size_t>::max()};
+ } else {
+ return *(*next_underline)++;
+ }
+}
+
+void PrintChunk(const std::pair<size_t, size_t>& lines_range,
+ const std::pair<UnderlineIter, UnderlineIter>& underlines, const std::string& text,
+ const std::vector<size_t>& line_starts, const DocPrinterOptions& options,
+ size_t line_number_width, std::string* out) {
+ UnderlineIter next_underline = underlines.first;
+ ByteSpan current_underline = PopNextUnderline(&next_underline, underlines.second);
+
+ for (size_t line_idx = lines_range.first; line_idx < lines_range.second; ++line_idx) {
+ if (options.print_line_numbers) {
+ std::string line_num_str = std::to_string(line_idx + 1);
+ line_num_str.push_back(' ');
+ for (size_t i = line_num_str.size(); i < line_number_width; ++i) {
+ out->push_back(' ');
+ }
+ *out += line_num_str;
+ }
+
+ size_t line_start = line_starts.at(line_idx);
+ size_t line_end =
+ line_idx + 1 == line_starts.size() ? text.size() : line_starts.at(line_idx + 1);
+ out->append(text.begin() + line_start, text.begin() + line_end);
+
+ bool printed_underline = false;
+ size_t line_pos = line_start;
+ bool printed_extra_caret = 0;
+ while (current_underline.first < line_end) {
+ if (!printed_underline) {
+ *out += std::string(line_number_width, ' ');
+ printed_underline = true;
+ }
+
+ size_t underline_end_for_line = std::min(line_end, current_underline.second);
+ size_t num_spaces = GetTextWidth(text, {line_pos, current_underline.first});
+ if (num_spaces > 0 && printed_extra_caret) {
+ num_spaces -= 1;
+ printed_extra_caret = false;
+ }
+ *out += std::string(num_spaces, ' ');
+
+ size_t num_carets = GetTextWidth(text, {current_underline.first, underline_end_for_line});
+ if (num_carets == 0 && !printed_extra_caret) {
+ // Special case: when underlineing an empty or unprintable string, make sure to print
+ // at least one caret still.
+ num_carets = 1;
+ printed_extra_caret = true;
+ } else if (num_carets > 0 && printed_extra_caret) {
+ num_carets -= 1;
+ printed_extra_caret = false;
+ }
+ *out += std::string(num_carets, '^');
+
+ line_pos = current_underline.first = underline_end_for_line;
+ if (current_underline.first == current_underline.second) {
+ current_underline = PopNextUnderline(&next_underline, underlines.second);
+ }
+ }
+
+ if (printed_underline) {
+ out->push_back('\n');
+ }
+ }
+}
+
+void PrintCut(size_t num_lines_skipped, std::string* out) {
+ if (num_lines_skipped != 0) {
+ std::ostringstream s;
+ s << "(... " << num_lines_skipped << " lines skipped ...)\n";
+ *out += s.str();
+ }
+}
+
+std::pair<size_t, size_t> GetLinesForUnderline(const ByteSpan& underline,
+ const std::vector<size_t>& line_starts,
+ size_t num_lines, const DocPrinterOptions& options) {
+ size_t first_line_of_underline = GetLineIndex(underline.first, line_starts);
+ size_t first_line_of_chunk = MoveBack(first_line_of_underline, options.num_context_lines);
+ size_t end_line_of_underline = GetLineIndex(underline.second - 1, line_starts) + 1;
+ size_t end_line_of_chunk =
+ MoveForward(end_line_of_underline, options.num_context_lines, num_lines);
+
+ return {first_line_of_chunk, end_line_of_chunk};
+}
+
+// If there is only one line between the chunks, it is better to print it as is,
+// rather than something like "(... 1 line skipped ...)".
+constexpr const size_t kMinLinesToCutOut = 2;
+
+bool TryMergeChunks(std::pair<size_t, size_t>* cur_chunk,
+ const std::pair<size_t, size_t>& new_chunk) {
+ if (new_chunk.first < cur_chunk->second + kMinLinesToCutOut) {
+ cur_chunk->second = new_chunk.second;
+ return true;
+ } else {
+ return false;
+ }
+}
+
+size_t GetNumLines(const std::string& text, const std::vector<size_t>& line_starts) {
+ if (line_starts.back() == text.size()) {
+ // Final empty line doesn't count as a line
+ return line_starts.size() - 1;
+ } else {
+ return line_starts.size();
+ }
+}
+
+size_t GetLineNumberWidth(size_t num_lines, const DocPrinterOptions& options) {
+ if (options.print_line_numbers) {
+ return std::to_string(num_lines).size() + 1;
+ } else {
+ return 0;
+ }
+}
+
+std::string DecorateText(const std::string& text, const std::vector<size_t>& line_starts,
+ const DocPrinterOptions& options,
+ const std::vector<ByteSpan>& underlines) {
+ size_t num_lines = GetNumLines(text, line_starts);
+ size_t line_number_width = GetLineNumberWidth(num_lines, options);
+
+ std::string ret;
+ if (underlines.empty()) {
+ PrintChunk({0, num_lines}, {underlines.begin(), underlines.begin()}, text, line_starts, options,
+ line_number_width, &ret);
+ return ret;
+ }
+
+ size_t last_end_line = 0;
+ std::pair<size_t, size_t> cur_chunk =
+ GetLinesForUnderline(underlines[0], line_starts, num_lines, options);
+ if (cur_chunk.first < kMinLinesToCutOut) {
+ cur_chunk.first = 0;
+ }
+
+ auto first_underline_in_cur_chunk = underlines.begin();
+ for (auto underline_it = underlines.begin() + 1; underline_it != underlines.end();
+ ++underline_it) {
+ std::pair<size_t, size_t> new_chunk =
+ GetLinesForUnderline(*underline_it, line_starts, num_lines, options);
+
+ if (!TryMergeChunks(&cur_chunk, new_chunk)) {
+ PrintCut(cur_chunk.first - last_end_line, &ret);
+ PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underline_it}, text, line_starts,
+ options, line_number_width, &ret);
+ last_end_line = cur_chunk.second;
+ cur_chunk = new_chunk;
+ first_underline_in_cur_chunk = underline_it;
+ }
+ }
+
+ PrintCut(cur_chunk.first - last_end_line, &ret);
+ if (num_lines - cur_chunk.second < kMinLinesToCutOut) {
+ cur_chunk.second = num_lines;
+ }
+ PrintChunk(cur_chunk, {first_underline_in_cur_chunk, underlines.end()}, text, line_starts,
+ options, line_number_width, &ret);
+ PrintCut(num_lines - cur_chunk.second, &ret);
+ return ret;
+}
+
+} // anonymous namespace
+
+DocPrinter::DocPrinter(const DocPrinterOptions& options) : options_(options) {
+ line_starts_.push_back(0);
+}
+
+void DocPrinter::Append(const Doc& doc) { Append(doc, NullOpt); }
+
+void DocPrinter::Append(const Doc& doc, Optional<ObjectPath> path_to_underline) {
+ path_to_underline_ = path_to_underline;
+ current_max_path_length_ = 0;
+ current_underline_candidates_.clear();
+ PrintDoc(doc);
+
+ underlines_.insert(underlines_.end(), current_underline_candidates_.begin(),
+ current_underline_candidates_.end());
+}
String DocPrinter::GetString() const {
std::string text = output_.str();
+
+ // Remove any trailing indentation
+ while (!text.empty() && text.back() == ' ') {
+ text.pop_back();
+ }
+
if (!text.empty() && text.back() != '\n') {
text.push_back('\n');
}
- return text;
+
+ std::vector<ByteSpan> underlines = underlines_;
+ SortAndMergeSpans(&underlines);
+ return DecorateText(text, line_starts_, options_, underlines);
}
void DocPrinter::PrintDoc(const Doc& doc) {
+ size_t start_pos = output_.tellp();
+
if (const auto* doc_node = doc.as<LiteralDocNode>()) {
PrintTypedDoc(GetRef<LiteralDoc>(doc_node));
} else if (const auto* doc_node = doc.as<IdDocNode>()) {
@@ -84,6 +321,24 @@ void DocPrinter::PrintDoc(const Doc& doc) {
LOG(FATAL) << "Do not know how to print " << doc->GetTypeKey();
throw;
}
+
+ size_t end_pos = output_.tellp();
+ for (const ObjectPath& path : doc->source_paths) {
+ MarkSpan({start_pos, end_pos}, path);
+ }
+}
+
+void DocPrinter::MarkSpan(const ByteSpan& span, const ObjectPath& path) {
+ if (path_to_underline_.defined()) {
+ if (path->Length() >= current_max_path_length_ &&
+ path->IsPrefixOf(path_to_underline_.value())) {
+ if (path->Length() > current_max_path_length_) {
+ current_max_path_length_ = path->Length();
+ current_underline_candidates_.clear();
+ }
+ current_underline_candidates_.push_back(span);
+ }
+ }
}
} // namespace printer
diff --git a/src/script/printer/base_doc_printer.h b/src/script/printer/base_doc_printer.h
index 8633dd0ded..f3fb24d946 100644
--- a/src/script/printer/base_doc_printer.h
+++ b/src/script/printer/base_doc_printer.h
@@ -22,14 +22,37 @@
#include <tvm/script/printer/doc.h>
#include <tvm/script/printer/doc_printer.h>
+#include <limits>
#include <memory>
#include <ostream>
#include <string>
+#include <utility>
+#include <vector>
namespace tvm {
namespace script {
namespace printer {
+/*! \brief Range of byte offsets in a string */
+using ByteSpan = std::pair<size_t, size_t>;
+
+/*! \brief Options to customize DocPrinter's output */
+struct DocPrinterOptions {
+ /*! \brief Number of spaces for one level of indentation */
+ int indent_spaces = 4;
+
+ /*! \brief Whether to print the line numbers */
+ bool print_line_numbers = false;
+
+ /*!
+ * \brief Number of context lines to print around the underlined text.
+ *
+ * If set to a non-default value `n`, only print `n` context lines before and after
+ * the underlined pieces of text.
+ */
+ size_t num_context_lines = std::numeric_limits<size_t>::max();
+};
+
/*!
* \brief DocPrinter is responsible for printing Doc tree into text format
* \details This is the base class for translating Doc into string.
@@ -45,7 +68,7 @@ class DocPrinter {
*
* \param options the option for printer
*/
- explicit DocPrinter(int indent_spaces = 4);
+ explicit DocPrinter(const DocPrinterOptions& options);
virtual ~DocPrinter() = default;
/*!
@@ -57,6 +80,16 @@ class DocPrinter {
*/
void Append(const Doc& doc);
+ /*!
+ * \brief Append a doc to the final content
+ *
+ * \param doc Doc to be printed
+ * \param path_to_underline Object path to be underlined
+ *
+ * \sa GetString
+ */
+ void Append(const Doc& doc, Optional<ObjectPath> path_to_underline);
+
/*!
* \brief Get the printed string of all Doc appended
*
@@ -192,13 +225,13 @@ class DocPrinter {
* \brief Increase the indent level of any content to be
* printed after this call
*/
- void IncreaseIndent() { indent_ += indent_spaces_; }
+ void IncreaseIndent() { indent_ += options_.indent_spaces; }
/*!
* \brief Decrease the indent level of any content to be
* printed after this call
*/
- void DecreaseIndent() { indent_ -= indent_spaces_; }
+ void DecreaseIndent() { indent_ -= options_.indent_spaces; }
/*!
* \brief Add a new line into the output stream
@@ -207,6 +240,7 @@ class DocPrinter {
*/
std::ostream& NewLine() {
output_ << "\n";
+ line_starts_.push_back(output_.tellp());
output_ << std::string(indent_, ' ');
return output_;
}
@@ -222,11 +256,31 @@ class DocPrinter {
std::ostringstream output_;
private:
- /*! \brief the number of spaces for one level of indentation */
- int indent_spaces_ = 4;
+ void MarkSpan(const ByteSpan& span, const ObjectPath& path);
+
+ /*! \brief Options to customize certain aspects of the output */
+ DocPrinterOptions options_;
/*! \brief the current level of indent */
int indent_ = 0;
+
+ /*! \brief For each line in the output_, byte offset of its first character */
+ std::vector<size_t> line_starts_;
+
+ /*! \brief Path of the object that we would like to underline */
+ Optional<ObjectPath> path_to_underline_;
+
+ /*!
+ * \brief Candidate spans to be underlined, until we find a better match.
+ * (A better match is an object with a longer path that is still a prefix of path_to_underline_.)
+ */
+ std::vector<ByteSpan> current_underline_candidates_;
+
+ /*! \brief Path length of the objects that are current candidates for underlining. */
+ int current_max_path_length_;
+
+ /*! \brief Spans that we have already committed to underline. */
+ std::vector<ByteSpan> underlines_;
};
} // namespace printer
diff --git a/src/script/printer/python_doc_printer.cc b/src/script/printer/python_doc_printer.cc
index 536c57abd9..d3a991d380 100644
--- a/src/script/printer/python_doc_printer.cc
+++ b/src/script/printer/python_doc_printer.cc
@@ -138,7 +138,7 @@ ExprPrecedence GetExprPrecedence(const ExprDoc& doc) {
class PythonDocPrinter : public DocPrinter {
public:
- explicit PythonDocPrinter(int indent_spaces = 4) : DocPrinter(indent_spaces) {}
+ explicit PythonDocPrinter(const DocPrinterOptions& options) : DocPrinter(options) {}
protected:
using DocPrinter::PrintDoc;
@@ -622,9 +622,17 @@ void PythonDocPrinter::PrintTypedDoc(const ClassDoc& doc) {
NewLineWithoutIndent();
}
-String DocToPythonScript(Doc doc, int indent_spaces) {
- PythonDocPrinter printer(indent_spaces);
- printer.Append(doc);
+String DocToPythonScript(Doc doc, int indent_spaces, bool print_line_numbers, int num_context_lines,
+ Optional<ObjectPath> path_to_underline) {
+ DocPrinterOptions options;
+ options.indent_spaces = indent_spaces;
+ options.print_line_numbers = print_line_numbers;
+ if (num_context_lines >= 0) {
+ options.num_context_lines = num_context_lines;
+ }
+
+ PythonDocPrinter printer(options);
+ printer.Append(doc, path_to_underline);
return printer.GetString();
}
diff --git a/tests/python/unittest/test_tvmscript_printer_underlining.py b/tests/python/unittest/test_tvmscript_printer_underlining.py
new file mode 100644
index 0000000000..a7e7dffb8b
--- /dev/null
+++ b/tests/python/unittest/test_tvmscript_printer_underlining.py
@@ -0,0 +1,361 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Optional
+
+import pytest
+
+from tvm.runtime import ObjectPath
+from tvm.script.printer.doc import (
+ StmtBlockDoc,
+ ExprStmtDoc,
+ IdDoc,
+ OperationDoc,
+ OperationKind,
+)
+from tvm.script.printer.doc_printer import to_python_script
+
+
+def make_path(name: str) -> ObjectPath:
+ return ObjectPath.root().attr(name)
+
+
+def make_id_doc(name: str, path_name: Optional[str] = None) -> IdDoc:
+ if path_name is None:
+ path_name = name
+ doc = IdDoc(name)
+ doc.source_paths = [make_path(path_name)]
+ return doc
+
+
+def format_script(s: str) -> str:
+ """
+ Remove leading and trailing blank lines, and make the minimum idention 0
+ """
+ s = s.strip("\n")
+
+ non_empty_lines = [line for line in s.splitlines() if line and not line.isspace()]
+ if not non_empty_lines:
+ # no actual content
+ return "\n"
+
+ line_indents = [len(line) - len(line.lstrip(" ")) for line in non_empty_lines]
+ spaces_to_remove = min(line_indents)
+
+ cleaned_lines = "\n".join(line[spaces_to_remove:] for line in s.splitlines())
+ if not cleaned_lines.endswith("\n"):
+ cleaned_lines += "\n"
+ return cleaned_lines
+
+
+def test_underline_basic():
+ doc = StmtBlockDoc(
+ [
+ ExprStmtDoc(make_id_doc("foo")),
+ ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("bar"), make_id_doc("baz")])),
+ ExprStmtDoc(make_id_doc("qux")),
+ ]
+ )
+ assert to_python_script(doc, path_to_underline=make_path("baz")) == format_script(
+ """
+ foo
+ bar + baz
+ ^^^
+ qux
+ """
+ )
+
+
+def test_underline_multiple_spans():
+ doc = StmtBlockDoc(
+ [
+ ExprStmtDoc(make_id_doc("foo")),
+ ExprStmtDoc(make_id_doc("bar")),
+ ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("foo"), make_id_doc("foo")])),
+ ]
+ )
+ assert to_python_script(doc, path_to_underline=make_path("foo")) == format_script(
+ """
+ foo
+ ^^^
+ bar
+ foo + foo
+ ^^^ ^^^
+ """
+ )
+
+
+def test_underline_multiple_spans_with_line_numbers():
+ doc = StmtBlockDoc(
+ [
+ ExprStmtDoc(make_id_doc("foo")),
+ ExprStmtDoc(make_id_doc("bar")),
+ ExprStmtDoc(OperationDoc(OperationKind.Add, [make_id_doc("foo"), make_id_doc("foo")])),
+ ]
+ )
+ assert to_python_script(
+ doc, print_line_numbers=True, path_to_underline=make_path("foo")
+ ) == format_script(
+ """
+ 1 foo
+ ^^^
+ 2 bar
+ 3 foo + foo
+ ^^^ ^^^
+ """
+ )
+
+
+def test_underline_multiline():
+ doc = StmtBlockDoc(
+ [
+ ExprStmtDoc(IdDoc("foo")),
+ ExprStmtDoc(IdDoc("bar")),
+ ]
+ )
+ doc.source_paths = [make_path("whole_doc")]
+
+ assert to_python_script(doc, path_to_underline=make_path("whole_doc")) == format_script(
+ """
+ foo
+ ^^^
+ bar
+ ^^^
+ """
+ )
+
+
+@pytest.mark.parametrize(
+ "to_underline, expected_text",
+ [
+ (
+ [0],
+ """
+ x0
+ ^^
+ x1
+ x2
+ (... 7 lines skipped ...)
+ """,
+ ),
+ (
+ [1],
+ """
+ x0
+ x1
+ ^^
+ x2
+ x3
+ (... 6 lines skipped ...)
+ """,
+ ),
+ (
+ [3],
+ """
+ x0
+ x1
+ x2
+ x3
+ ^^
+ x4
+ x5
+ (... 4 lines skipped ...)
+ """,
+ ),
+ (
+ [4],
+ """
+ (... 2 lines skipped ...)
+ x2
+ x3
+ x4
+ ^^
+ x5
+ x6
+ (... 3 lines skipped ...)
+ """,
+ ),
+ (
+ [6],
+ """
+ (... 4 lines skipped ...)
+ x4
+ x5
+ x6
+ ^^
+ x7
+ x8
+ x9
+ """,
+ ),
+ (
+ [9],
+ """
+ (... 7 lines skipped ...)
+ x7
+ x8
+ x9
+ ^^
+ """,
+ ),
+ (
+ [0, 9],
+ """
+ x0
+ ^^
+ x1
+ x2
+ (... 4 lines skipped ...)
+ x7
+ x8
+ x9
+ ^^
+ """,
+ ),
+ (
+ [0, 3, 9],
+ """
+ x0
+ ^^
+ x1
+ x2
+ x3
+ ^^
+ x4
+ x5
+ x6
+ x7
+ x8
+ x9
+ ^^
+ """,
+ ),
+ (
+ [0, 6, 9],
+ """
+ x0
+ ^^
+ x1
+ x2
+ x3
+ x4
+ x5
+ x6
+ ^^
+ x7
+ x8
+ x9
+ ^^
+ """,
+ ),
+ (
+ [33],
+ """
+ x0
+ x1
+ x2
+ x3
+ x4
+ x5
+ x6
+ x7
+ x8
+ x9
+ """,
+ ),
+ ],
+)
+def test_print_two_context_lines(to_underline, expected_text):
+ doc = StmtBlockDoc(
+ [ExprStmtDoc(make_id_doc(f"x{i}", "yes" if i in to_underline else "no")) for i in range(10)]
+ )
+ result = to_python_script(doc, num_context_lines=2, path_to_underline=make_path("yes"))
+ assert result == format_script(expected_text)
+
+
+def test_underline_and_print_line_numbers():
+ doc = StmtBlockDoc([ExprStmtDoc(make_id_doc(f"line{i + 1}")) for i in range(12)])
+ result = to_python_script(doc, print_line_numbers=True, path_to_underline=make_path("line6"))
+ assert result == format_script(
+ """
+ 1 line1
+ 2 line2
+ 3 line3
+ 4 line4
+ 5 line5
+ 6 line6
+ ^^^^^
+ 7 line7
+ 8 line8
+ 9 line9
+ 10 line10
+ 11 line11
+ 12 line12
+ """
+ )
+
+
+def test_underline_and_print_line_numbers_with_context():
+ doc = StmtBlockDoc([ExprStmtDoc(make_id_doc(f"line{i + 1}")) for i in range(12)])
+ result = to_python_script(
+ doc, print_line_numbers=True, num_context_lines=2, path_to_underline=make_path("line8")
+ )
+ assert result == format_script(
+ """
+ (... 5 lines skipped ...)
+ 6 line6
+ 7 line7
+ 8 line8
+ ^^^^^
+ 9 line9
+ 10 line10
+ (... 2 lines skipped ...)
+ """
+ )
+
+
+def test_underline_based_on_path_prefix():
+ doc = StmtBlockDoc([ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc(make_id_doc("bar"))])
+ result = to_python_script(doc, path_to_underline=make_path("foo").attr("x").attr("y"))
+ # There is no document that matches the desired path exactly,
+ # but path of "foo" is a prefix of the desired path, and thus should be underlined.
+ assert result == format_script(
+ """
+ foo
+ ^^^
+ bar
+ """
+ )
+
+
+def test_longer_prefix_must_win():
+ foo_x = IdDoc("foo_x")
+ foo_x.source_paths = [make_path("foo").attr("x")]
+
+ doc = StmtBlockDoc(
+ [ExprStmtDoc(make_id_doc("foo")), ExprStmtDoc(make_id_doc("bar")), ExprStmtDoc(foo_x)]
+ )
+ result = to_python_script(doc, path_to_underline=make_path("foo").attr("x").attr("y"))
+ # "foo" should not be underlined because there is a document with a more specific path prefix
+ assert result == format_script(
+ """
+ foo
+ bar
+ foo_x
+ ^^^^^
+ """
+ )