『065:文字列を最大nバイトに切り詰める』は codecs モジュールまかせでもできた

終わった…あっさりと。文字コードと戯れていたオレってなんだったんだろう。

# coding: utf-8

from codecs import getincrementaldecoder, iterencode

def _iter_mbchar(iterator, encoding, errors='strict', **kwargs):
    decoder = getincrementaldecoder(encoding)(errors, **kwargs)
    mbchar = []
    for input in iterator:
        mbchar.append(input)
        output = decoder.decode(input)
        if output:
            yield ''.join(mbchar)
            del mbchar[:]
    output = decoder.decode("", True)
    if output:
        yield ''.join(mbchar)

def take_nbytes_str(str_, n, encoding, errors='strict', **kwargs):
    u"""
    文字列 str_ を最大 n バイトに切り詰める

    str_:     8 ビット文字列
    n:        最大のバイト数。これより長い場合、切り詰める
    encoding: str_ の文字コード
    errors:   コーデックが用いるエラー処理方法
    **kwargs: コーデックに渡すオプション
    """
    length = 0
    temp = []
    for mbchar in _iter_mbchar(str_, encoding, errors, **kwargs):
        length += len(mbchar)
        if length > n:
            break
        temp.append(mbchar)
    return ''.join(temp)

def take_nbytes_unicode(unicode_, n, encoding, errors='strict', **kwargs):
    u"""
    ユニコード文字列 unicode_ をエンコードする、ただし最大 n バイトまで

    unicode_: ユニコード文字列
    n:        最大のバイト数。出力結果がこれより長くなることはない
    encoding: 文字コード
    errors:   コーデックが用いるエラー処理方法
    **kwargs: コーデックに渡すオプション
    """
    length = 0
    temp = []
    for mbchar in iterencode(unicode_, encoding, errors, **kwargs):
        length += len(mbchar)
        if length > n:
            break
        temp.append(mbchar)
    return ''.join(temp)

def _test():
    def _write(f, code):
        s = u'あいうえおかきくけこさしすせそたちつてと'.encode(code)
        for i in range(1, 40):
            f.write(''.join(take_nbytes_str(s, i, code)))
            f.write('\n')
        u = u'あいうえおかきくけこさしすせそたちつてと'
        for i in range(1, 40):
            f.write(''.join(take_nbytes_unicode(u, i, code)))
            f.write('\n')
    f_sjis = open('_sjis.txt', 'w')
    try:
        _write(f_sjis, 'shift_jis')
    finally:
        f_sjis.close()
    f_euc = open('_euc.txt', 'w')
    try:
        _write(f_euc, 'euc_jp')
    finally:
        f_euc.close()
    f_utf8 = open('_utf8.txt', 'w')
    try:
        _write(f_utf8, 'utf_8')
    finally:
        f_utf8.close()

if __name__ == '__main__':
    _test()