[GitHub] [tvm] junrushao commented on a diff in pull request #13839: [TVMScript] Comments and docstrings printing

["junrushao (via GitHub)" <gi...@apache.org>]

junrushao commented on code in PR #13839:
URL: https://github.com/apache/tvm/pull/13839#discussion_r1086976751


##########
src/script/printer/doc_printer/python_doc_printer.cc:
##########
@@ -642,11 +652,40 @@ void PythonDocPrinter::PrintTypedDoc(const ClassDoc& doc) {
   NewLineWithoutIndent();
 }
 
+void PythonDocPrinter::PrintTypedDoc(const CommentDoc& doc) {
+  if (doc->comment.defined()) {
+    MaybePrintCommenMultiLines(doc, false);
+  }
+}
+
+void PythonDocPrinter::PrintTypedDoc(const DocStringDoc& doc) {
+  if (doc->comment.defined() && !doc->comment.value().empty()) {
+    output_ << "\"\"\"" << doc->comment.value() << "\"\"\"";
+  }
+}
+
 String DocToPythonScript(Doc doc, const PrinterConfig& cfg) {
   if (cfg->num_context_lines < 0) {
     cfg->num_context_lines = std::numeric_limits<int32_t>::max();
   }
   PythonDocPrinter printer(cfg);
+  if (cfg->print_headers) {
+    Array<StmtDoc> headers;
+    if (cfg->ir_usage.count("ir")) {
+      headers.push_back(CommentDoc("from tvm.script import ir as " + cfg->ir_prefix));
+    }
+    if (cfg->ir_usage.count("tir")) {
+      headers.push_back(CommentDoc("from tvm.script import tir as " + cfg->tir_prefix));
+    }
+    if (cfg->ir_usage.count("relax")) {
+      headers.push_back(CommentDoc("from tvm.script import relax as " + cfg->relax_prefix));
+    }
+    if (headers.size()) {
+      // Add a blank line between headers and script codes
+      headers.push_back(CommentDoc(""));
+    }
+    printer.Append(StmtBlockDoc(headers));
+  }

Review Comment:
   We should move this logic to the caller site :-)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org