Pythonのassertで複数行を入出力させたいとき

お世話になっております。

競技プログラミングのような問題を解くことをPython3.6 + Pycharmで想定していて、簡単なテストをコード上で行いたいです。なのですが、複数行の入出力をassertさせる方法がわかりません。

複数行のinputやoutputをpythonでassertする方法(もしくはより簡単にPycharm上で入出力のテストができるやりかた)を教えていただけないでしょうか。

例えば、yukicoderのこの問題のサンプルとして

#input
5

#output
55555
5555
555
55
5

がありますが、以下のように書いてもassertでエラーになってしまいます。提出はできたのでmult_number関数の中身は問題ないと思われます。
(入出力はstrでされるものだと認識しています)

#num = int(input())

def mult_number(num):
    for i in range(0,num):
        print(str(num) * (num - i))

nums_4 = """
4444
444
44
4
"""

nums_5 = """
55555
5555
555
55
5
"""

assert mult_number(4) == nums_4
assert mult_number(5) == nums_5

Traceback (most recent call last):
4444
444
  File "input.py", line 22, in <module>
44
4
    assert mult_number(int(4)) == nums_4
AssertionError

Process finished with exit code 1

よろしくお願いいたします。

追記分

#関数にしてreturnにした。WAでテストケースで引っかかった
num = int(input())

def mult_number(num):
    for i in range(0,num):
        return(str(num) * (num - i))

#関数にしてprintにした。WAでテストケースで引っかかった
num = int(input())

def mult_number(num):
    for i in range(0,num):
        return(str(num) * (num - i))

#関数にしていない。　テストケースをすべて通過した
num = int(input())
for i in range(0,num):
    print(str(num) * (num - i))

行動規範の内容に同意します

回答2件

print(...)の内容をテストしたければ以下のようにします。

python
1from contextlib import redirect_stdout
2from io import StringIO
3
4
5fp = StringIO()
6with redirect_stdout(fp):
7    mult_number(4)
8assert fp.getvalue() == nums_4

投稿2018/02/19 11:54

YouheiSakurai

総合スコア6149

mofu_mofu

2018/02/19 12:30

YouheiSakurai様そんなやりかたがあるのですね。ご丁寧にありがとうございます。

行動規範の内容に同意します

ベストアンサー

nums_4は文字列ですが、mult_number(4)はNoneです。※
対話環境で同じように表示されるからと言って同値だとは全く言えないです。

Python
1>>> "aaa"
2'aaa'
3>>> print("aaa")
4aaa
5>>>
6>>> type("aaa")
7<class 'str'>
8>>> type(print("aaa"))
9aaa
10<class 'NoneType'>
11>>>
12>>> "aaa" == print("aaa")
13aaa
14False
15>>> assert "aaa" == print("aaa")
16aaa
17Traceback (most recent call last):
18  File "<stdin>", line 1, in <module>
19AssertionError

複数行の入出力をassertさせる方法がわかりません。

この場合行数は関係ないです。

※ Pythonでは、『戻り値なしの関数の戻り値』はNoneになります。ややこしい。

コメントを受けて

multi_numberを文字列を返す関数にしてみましょう。

リストを返す場合

Python
1def multi_number(num):
2    ret = []
3    for i in range(num):
4        ret.append(str(num) * (num-i))
5    return ret
6
7print(*multi_number(3), sep='\n')
8# 普通にfor文でぶん回しても良い
9# for row in multi_number(3):
10#     print(row)

ジェネレータ関数にする場合

Python
1def multi_number(num):
2    for i in range(num):
3        yield str(num) * (num-i)
4
5for row in multi_number(3):
6    print(row)

慣れるまでは前者を使った方が分かりやすいかもしれませんね。

忘れてた。
真の意味で、文字列を返す関数

Python
1def multi_number(num):
2    ret = ''
3    for i in range(num):
4        ret += str(num) * (num-i) + '\n'
5    return ret[:-1]
6
7print(multi_number(3))

ret[:-1]したのは最後に改行が重複するのを防ぐためです。
retのまま返して、print関数の引数endに空文字列を指定するのもアリです。

投稿2018/02/19 11:35

編集2018/02/19 12:37

LouiS0616

総合スコア35676

mofu_mofu

2018/02/19 12:07

なるほど。print()だとNoneになってreturn("hoge")だとstrになるのですね。追記したのですが、それとは別に関数ではないものを関数にして、それをテストしたのですが、関数でないものを関数にするときに挙動が変わっているようです。なのでこちらもおかしい原因になっているようです。 def mult_number(num): for i in range(0,num): return(str(num) * (num - i)) print(mult_number(3)) だと返り値が 333でテストケースが通らなかったです。