Python - Reading and Writing Files - PDB、Molecules, Serial Numbers

Practical Programming
楽天ブックス 紀伊国屋書店

開発環境

• macOS Catalina - Apple (OS)
• Emacs (Text Editor)
• Windows 10 Pro (OS)
• Visual Studio Code (Text Editor)
• Python 3.7 (プログラミング言語)

Practical Programming: An Introduction to Computer Science Using Python 3.6 (Paul Gries(著)、Jennifer Campbell(著)、Jason Montojo(著)、Pragmatic Bookshelf)のChapter 10(Reading and Writing Files)、Exercise 8の解答を求めてみる。

コード

#!/usr/bin/env python3
from unittest import TestCase, main
from test.support import captured_stdout
from typing import TextIO, Union
from io import StringIO

print('8.')


class TestReadAllMolecules(TestCase):
    def test(self):
        compound1 = '''COMPND T1
ATOM 1 N 0.1 0.2 0.3
ATOM 1 N 0.2 0.1 0.0
END
'''
        compound2 = '''COMPND T2
ATOM 1 A 0.1 0.2 0.3
ATOM 2 A 0.2 0.1 0.0
END
'''
        comment = "CMNT ...\n"
        blank_line = ' \n'
        in_file = StringIO(compound1 + comment +
                           blank_line + compound2 + blank_line + blank_line + comment)
        with captured_stdout() as stdout:
            result = read_all_molecules(in_file)
        self.assertEqual(stdout.getvalue(),
                         'serial number: got 1, expected 2\n')
        self.assertEqual(len(result), 2)
        self.assertEqual(result[0], ['T1', ['N', '0.1', '0.2', '0.3'], [
                         'N', '0.2', '0.1', '0.0']])
        self.assertEqual(result[1], ['T2', ['A', '0.1', '0.2', '0.3'], [
                         'A', '0.2', '0.1', '0.0']])


def skip(line: str) -> bool:
    return line.startswith('CMNT') or line.isspace()


def read_molecule(reader: TextIO) -> Union[list, None]:
    while True:
        line = reader.readline()
        if not line:
            return None
        if skip(line):
            continue
        break
    parts = line.split()
    name = parts[1]
    molecule: list = [name]
    serial_number = 1
    for line in reader:
        if skip(line):
            continue
        if line.startswith('END'):
            break
        parts = line.split()
        number = int(parts[1])
        if number != serial_number:
            print(f'serial number: got {number}, expected {serial_number}')
        molecule.append(parts[2:])
        serial_number += 1
    return molecule


def read_all_molecules(reader: TextIO) -> list:
    result = []
    while True:
        molecule = read_molecule(reader)
        if molecule is None:
            break
        result.append(molecule)
    return result


if __name__ == "__main__":
    main()

入出力結果(Zsh、PowerShell、Terminal、Jupyter(IPython))

% ./sample8.py -v
8.
test (__main__.TestReadAllMolecules) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.000s

OK
%