diff --git a/example.cpp b/example.cpp index 5acaf1b..bd88056 100644 --- a/example.cpp +++ b/example.cpp @@ -27,6 +27,9 @@ signed main() { a.foo(); a.bar(3,4,"hi"); + print("this should not be in the output file"); + redirect_output("example.out"); + auto [i, j] = read(); print(i, j); @@ -68,4 +71,12 @@ signed main() { print(map{{3,"three"},{1,"one"}}); print(tuple{unordered_map{{3,"three"},{1,"one"}}}); print(unordered_set{3,1,4}); + + redirect_input("example.in"); + print("reading from example.in:", read_string()); + redirect_input("example.in"); + print("reading from example.in:", read_string()); + print("this is the last line on example.out"); + reset_output(); + print("this is stdout only"); } diff --git a/include/bits/soi-redirect.hpp b/include/bits/soi-redirect.hpp new file mode 100644 index 0000000..93a7b95 --- /dev/null +++ b/include/bits/soi-redirect.hpp @@ -0,0 +1,93 @@ +#include +#include + +namespace soi { + +namespace redirect { + +class teebuf : public std::streambuf { + std::streambuf *sb1_; + std::streambuf *sb2_; + + int overflow(int c) { + using traits = std::streambuf::traits_type; + bool rc = true; + if (!traits::eq_int_type(traits::eof(), c)) { + traits::eq_int_type(this->sb1_->sputc(c), traits::eof()) && (rc = false); + traits::eq_int_type(this->sb2_->sputc(c), traits::eof()) && (rc = false); + } + return rc ? traits::not_eof(c) : traits::eof(); + } + int sync() { + bool rc = false; + this->sb1_->pubsync() != -1 || (rc = false); + this->sb2_->pubsync() != -1 || (rc = false); + return rc ? -1 : 0; + } + +public: + teebuf(std::streambuf *sb1, std::streambuf *sb2) : sb1_(sb1), sb2_(sb2) {} +}; + +static std::streambuf *checked_rdbuf(std::ios &stream, const char *filename, + const char *action) { + if (!stream) { + std::cerr << "ERROR: can't open to file " << filename << " for " << action + << '\n'; + std::exit(-1); + } + return stream.rdbuf(); +} + +struct output_redirector { + std::ofstream fout; + teebuf tee; + std::streambuf *old_cout; + + output_redirector(const char *filename) + : fout(filename), + tee(checked_rdbuf(fout, filename, "writing"), std::cout.rdbuf()), + old_cout(std::cout.rdbuf(&tee)) {} + + ~output_redirector() { std::cout.rdbuf(old_cout); } +}; + +struct input_redirector { + std::ifstream fin; + std::streambuf *old_cin; + + input_redirector(const char *filename) + : fin(filename), + old_cin(std::cin.rdbuf(checked_rdbuf(fin, filename, "reading"))) {} + + ~input_redirector() { std::cin.rdbuf(old_cin); } +}; + +std::forward_list output_redirectors; +std::unique_ptr the_input_redirector; + +} // namespace redirect + +} // namespace soi + +void redirect_output(const char *filename) { + soi::redirect::output_redirectors.emplace_front(filename); +} +void redirect_output(std::string const &filename) { + redirect_output(filename.c_str()); +} +void reset_output() { + soi::redirect::output_redirectors.clear(); +} + +void reset_input() { + soi::redirect::the_input_redirector.reset(); +} +void redirect_input(const char *filename) { + reset_input(); + auto next = new soi::redirect::input_redirector(filename); + soi::redirect::the_input_redirector.reset(next); +} +void redirect_input(std::string const &filename) { + redirect_input(filename.c_str()); +} diff --git a/include/soi b/include/soi index aa7518b..e1265bd 100644 --- a/include/soi +++ b/include/soi @@ -77,7 +77,7 @@ soi_initializer soi_initializer_{false}; } // end namespace soi #include "bits/soi-io.hpp" - +#include "bits/soi-redirect.hpp" #include "bits/soi-deprecate.hpp" #define int int64_t