/dev/null with pybind11 and unittest.mock

develop
, , ,

嗨嗨，我愛 pybind11 ！

Introduction

最近工作上用 pybind11 來接自己包出來會 Link 到 TensorFlow 的 C++ Library。覺得相當好用寫起來又乾淨，還可以傳入傳出 NumPy 的 ndarray， 實在是太方便了啊。

另外也使用 Python 的 unittest 工具來寫許多的測試， 反正有新功能要加就也會先寫些測試項目，另外也會驗證會不會影響到原本的功能。 通常跑 unittest 的時候，除非要看特定跑過哪些項目， 不然其實是希望只要印出跑完測試的結果就好，像是底下這樣，

1
2
3
4
5
6
$ make tests
.......s...........
----------------------------------------------------------------------
Ran 19 tests in 0.002s

OK (skipped=1)

通常在可以控制的範圍下，都是沒問題的，像是有些 Warnings 只要修一下， 就不會影響到上面乾淨的輸出結果，

1
2
3
with warnings.catch_warnings():
    warnings.simplefilter("ignore", category=PendingDeprecationWarning)
    import np.tools.tflite_runner

不過如剛剛有提到，有些其實是呼叫到別人的 Library ， 尤其是一些提示錯誤的訊息像是檔案找不到之類的，本來就會想要從 stderr 印出來， 在故意寫這種的測試項目時，其實就會看到錯誤訊息截斷了測試結果。

1
2
3
4
5
6
7
8
python tests/mock_stdout_tests.py TestFooBarStdout.test_print_printf TestFooBarStdout.test_print_fprintf_stdout
printf: [test_print_printf]
.fprintf_stdout: [test_print_fprintf_stdout]
.
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

所以這篇會以 Python unittest 的角度，分別將 Python 的 print，C++ 的 cout/cerr 以及 C 的 printf/fprintf 的結果導走，不讓這些訊息影響到測試的結果。

print in Python with unittest.mock

unittest.mock 是個可以在 Python 的 unittest 環境下， 用假造的物件替換掉原本應該真正使用的部份，其實功能相當多，這邊只用來替換 sys.stdout 就可以將 print() 的訊息截斷並拉回來，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class TestStdout(unittest.TestCase):
    @unittest.skip('skip due to annoying stdout print')
    def test_print_stdout(self):
        print('test_print_stdout')

    @unittest.skip('skip due to annoying stderr print')
    def test_print_stderr(self):
        print('test_print_stderr', file=sys.stderr)


class TestMockStdout(unittest.TestCase):
    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_stdout(self, mock_stdout):
        print('test_print_mock_stdout')

    @unittest.mock.patch('sys.stderr', new_callable=io.StringIO)
    def test_print_mock_stderr(self, mock_stderr):
        print('test_print_mock_stderr', file=sys.stderr)

如上面的測試項目，test_print_mock_stderr 就不會像 test_print_stderr 一樣印出訊息，因為被 mock_stderr 給接過來了。如果都是純的Python 程式， 這樣就非常足夠且方便使用。

cout/cerr in C++ with py::scoped_estream_redirect in pybind11

python module with cout

不過要是底下有些 python 的 module 是用 C++ 接的，然後呼叫 std::cout/std::cerr 印出訊息，上面的方法還不夠，為了驗證首先就用 pybind11 寫一個簡單的 module，

1
2
3
4
5
6
7
8
9
10
#include <pybind11/pybind11.h>
#include <pybind11/iostream.h>

namespace py = pybind11;

PYBIND11_MODULE(print, m) {
    m.attr("__name__") = "foo.bar.print";
    m.def("cout", [](const char* str) {
            std::cout << "cout: [" << str << "]" << std::endl; });
}

也寫個 unittest 來測試一下，結果還真的 mock 不起來， 所以把它 skip 掉了。

1
2
3
4
5
6
class TestFooBarMockStdout(unittest.TestCase):
    @unittest.skip('skip due to annoying stdout cout even mocked')
    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_cout(self, mock_stdout):
        foo.bar.print.cout('test_print_mock_cout')

py::scoped_ostream_redirect in pybind11

其實 pybind11 的 Capturing standard output from ostream， 就有提到這個問題，

Often, a library will use the streams std::cout and std::cerr to print, but this does not play well with Python’s standard sys.stdout and sys.stderr redirection. Replacing a library’s printing with py::print may not be feasible.

所以 pybind11 也有提供 py::scoped_ostream_redirect 的方法， 將 std::cout 重導到 Python 的 sys.stdout 去，用 Capturing standard output from ostream，改寫上面的範例，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include <pybind11/pybind11.h>
#include <pybind11/iostream.h>

namespace py = pybind11;

PYBIND11_MODULE(print, m) {
    m.attr("__name__") = "foo.bar.print";
    m.def("cout", [](const char* str) {
            std::cout << "cout: [" << str << "]" << std::endl; })
     .def("cout_redirect", [](const char* str) {
            std::cout << "cout redirect: [" << str << "]" << std::endl; },
        py::call_guard<py::scoped_ostream_redirect,
                       py::scoped_estream_redirect>());
}
1
2
3
4
5
6
7
8
9
class TestFooBarMockStdout(unittest.TestCase):
    @unittest.skip('skip due to annoying stdout cout even mocked')
    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_cout(self, mock_stdout):
        foo.bar.print.cout('test_print_mock_cout')

    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_cout_redirect(self, mock_stdout):
        foo.bar.print.cout_redirect('test_print_mock_cout_redirect')

呼叫有被 guard 的 cout_redirect，就可以順利地被 mock_stdout 接走了， 當然如果沒有使用 unittest.mock**，訊息一樣是會顯示在 **sys.stdout 上。

printf/fprintf in C with os.dup/os.dup2

但如果不是使用 std::cout/std::cerr ，而是使用 printf 或 fprintf 的話，py::scoped_ostream_redirect 也沒有辦法 Capturing standard output from ostream 也有說明，

The above methods will not redirect C-level output to file descriptors, such as fprintf. For those cases, you’ll need to redirect the file descriptors either directly in C or with Python’s os.dup2 function in an operating-system dependent way.

建議可以使用 os.dup2 自己接走。 既然如此就自己來實作 > /dev/null 2> /dev/null 的功能吧， 也可以達到不影響 unittest 結果顯示。首先還是來準備一下用 fprint() 的 Module，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include <pybind11/pybind11.h>
#include <pybind11/iostream.h>

namespace py = pybind11;

PYBIND11_MODULE(print, m) {
    m.attr("__name__") = "foo.bar.print";
    m.def("printf", [](const char* str){
            printf("printf: [%s]\n", str); })
     .def("fprintf_stdout", [](const char* str) {
             fprintf(stdout, "fprintf_stdout: [%s]\n", str); })
     .def("fprintf_stderr", [](const char* str) {
             fprintf(stderr, "fprintf_stderr: [%s]\n", str); })
     .def("fflush_stdout", []() { fflush(stdout); })
     .def("fflush_stderr", []() { fflush(stderr); });
}

果然即便是用了 unittest.mock ，還是沒辦法擷取到，只好 unittest.skip。

1
2
3
4
5
6
7
8
9
10
11
class TestFooBarMockStdout(unittest.TestCase):
    @unittest.skip('skip due to annoying stdout printf even mocked')
    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_printf(self, mock_stdout):
        foo.bar.print.printf('test_print_mock_printf')

    @unittest.skip('skip due to annoying stdout fprintf even mocked')
    @unittest.mock.patch('sys.stdout', new_callable=io.StringIO)
    def test_print_mock_fprintf_stdout(self, mock_stdout):
        foo.bar.print.fprintf_stdout('test_print_mock_fprintf_stdout')

所以來開個 /dev/null 用 os.dup/os.dup2 自己做重導吧， 然後跑完之後再恢復回來，該正常顯示的還是要可以顯示，

1
2
3
4
5
6
7
8
9
class TestFooBarStdout(unittest.TestCase):
    @unittest.skip('skip due to annoying stdout cout at the end')
    def test_dup2_print_printf(self):
        with open(os.devnull, 'w') as null:
            _dup = os.dup(1)
            os.dup2(null.fileno(), 1)
            foo.bar.print.printf('test_dup2_print_printf')
            null.flush()
            os.dup2(_dup, 1)

不過這樣還是有點問題，因為導到 /dev/null 的 stdout ， 雖然沒有馬上印在螢怒上，但還是會在程式結束之後再刷出來， 因為這時候為了其他正常的顯示需要，會將 stdout 接回來， 所以最後還是會在程式結束之後印，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
$ python tests/mock_stdout_tests.py TestFooBarStdout.test_print_fprintf_stdout
fprintf_stdout: [test_print_fprintf_stdout]
.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

$ python tests/mock_stdout_tests.py TestFooBarStdout.test_dup2_print_fprintf_stdout
.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK
fprintf_stdout: [test_dup2_print_fprintf_stdout]

這時候可以利用之前加的 fflush_stdout()/fflush_stderr()**， 在切換 **stdout/stderr 的時候強制刷一下，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class TestFooBarStdout(unittest.TestCase):
    def test_dup2_print_printf_and_fflush_stdout(self):
        with open(os.devnull, 'w') as null:
            _dup = os.dup(1)
            os.dup2(null.fileno(), 1)
            foo.bar.print.printf('test_dup2_print_printf_and_fflush_stdout')
            foo.bar.print.fflush_stdout()
            null.flush()
            os.dup2(_dup, 1)

    def test_dup2_print_fprintf_stdout_and_fflush_stdout(self):
        with open(os.devnull, 'w') as null:
            _dup = os.dup(1)
            os.dup2(null.fileno(), 1)
            foo.bar.print.fprintf_stdout('test_dup2_print_fprintf_stdout_and_fflush_stdout')
            foo.bar.print.fflush_stdout()
            null.flush()
            os.dup2(_dup, 1)

    def test_dup2_print_fprintf_stderr_and_fflush_stderr(self):
        with open(os.devnull, 'w') as null:
            _dup = os.dup(2)
            os.dup2(null.fileno(), 2)
            foo.bar.print.fprintf_stderr('test_dup2_print_fprintf_stderr_and_fflush_stderr')
            foo.bar.print.fflush_stderr()
            null.flush()
            os.dup2(_dup, 2)

如此一來就可以達到濾掉 printf()/fprintf() 的訊息，不會影響到結果顯示了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
$ python tests/mock_stdout_tests.py TestFooBarStdout.test_dup2_print_printf_and_fflush_stdout
.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

$ python tests/mock_stdout_tests.py TestFooBarStdout.test_dup2_print_fprintf_stdout_and_fflush_stdout
.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

$ python tests/mock_stdout_tests.py TestFooBarStdout.test_dup2_print_fprintf_stderr_and_fflush_stderr
.
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK