1 | #pragma once |
2 | |
3 | #include <sstream> |
4 | #include <iostream> |
5 | |
6 | namespace taichi { |
7 | |
8 | // PythonPrintBuffer holds the logs printed from kernel before sending them back |
9 | // to python. The name could be a bit misleading, as it is really just a string |
10 | // buffer, and can be used without Python. |
11 | struct PythonPrintBuffer { |
12 | std::stringstream ss; |
13 | bool enabled{false}; |
14 | |
15 | template <typename T> |
16 | PythonPrintBuffer &operator<<(const T &t) { |
17 | if (enabled) |
18 | ss << t; |
19 | else |
20 | std::cout << t; |
21 | return *this; |
22 | } |
23 | std::string pop_content() { |
24 | auto ret = ss.str(); |
25 | ss = std::stringstream(); |
26 | return ret; |
27 | } |
28 | }; |
29 | |
30 | extern PythonPrintBuffer py_cout; |
31 | |
32 | } // namespace taichi |
33 | |