source_location.h

#pragma once

#include <ostream>
#include <sstream>
#include <stdexcept>
#include <string>

namespace torch {
namespace jit {
// SourceLocation represents source code-level debug information for a node.
// It contains information about where a node got generated.
// In the case of tracing this will be a python stack trace.
// In the case of using the scripting frontend this will be backed
// by a SourceRange object
struct SourceLocation {
  virtual ~SourceLocation() = default;
  virtual void highlight(std::ostream& out) const = 0;

  std::string wrapException(
      const std::exception& e,
      const std::string& additional = "") {
    std::stringstream msg;
    msg << "\n" << e.what() << ":\n";
    if (!additional.empty()) {
      msg << additional << ":\n";
    }
    highlight(msg);
    return msg.str();
  }
  void wrapAndRethrowException(
      const std::exception& e,
      const std::string& additional = "") {
    throw std::runtime_error(wrapException(e, additional));
  }
};

inline std::ostream& operator<<(std::ostream& out, const SourceLocation& sl) {
  sl.highlight(out);
  return out;
}

// normally a python stack trace
struct StringSourceLocation : public SourceLocation {
  StringSourceLocation(std::string context) : context(std::move(context)) {}
  void highlight(std::ostream& out) const override {
    out << context;
  }

 private:
  std::string context;
};

} // namespace jit
} // namespace torch