For each built-in type annotation such as int str, list, set, tuple, and dict you can specify the kind of thing it contains. To explore this, let’s revisit atoms and molecules from , using dictionaries and tuples in addition to lists.
Recall functions read_molecule and read_all_molecules; here are the headers and docstrings:
| | def read_molecule(reader: TextIO) -> list: |
| | """Read a single molecule from reader and return it, or return None to |
| | signal end of file. The first item in the result is the name of the |
| | compound; each list contains an atom type and the X, Y, and Z coordinates |
| | of that atom. |
| | |
| | >>> instring = ('COMPND TEST\\nATOM 1 N 0.1 0.2 0.3\\n' + |
| | 'ATOM 2 N 0.2 0.1 0.0\\nEND\\n') |
| | >>> infile = StringIO(instring) |
| | >>> read_molecule(infile) |
| | ['TEST', ['N', '0.1', '0.2', '0.3'], ['N', '0.2', '0.1', '0.0']] |
| | """ |
| | def read_all_molecules(reader: TextIO) -> list: |
| | """Read zero or more molecules from reader, returning a list of the |
| | molecule information. |
| | |
| | >>> cmpnd1 = ('COMPND T1\\nATOM 1 N 0.1 0.2 0.3\\n' + |
| | 'ATOM 2 N 0.2 0.1 0.0\\nEND\\n') |
| | >>> cmpnd2 = ('COMPND T2\\nATOM 1 A 0.1 0.2 0.3\\n' + |
| | 'ATOM 2 A 0.2 0.1 0.0\\nEND\\n') |
| | >>> infile = StringIO(cmpnd1 + cmpnd2) |
| | >>> result = read_all_molecules(infile) |
| | >>> result[0] |
| | ['T1', ['N', '0.1', '0.2', '0.3'], ['N', '0.2', '0.1', '0.0']] |
| | >>> result[1] |
| | ['T2', ['A', '0.1', '0.2', '0.3'], ['A', '0.2', '0.1', '0.0']] |
| | """ |
Assuming that molecules have unique names, it would make sense for read_all_molecules to return a dictionary where the keys are the names of compounds and the values are the atoms.
Additionally, instead of using a four-item list for atoms, each atom will be represented as a tuple where the first item is the type of the atom and the second item is a tuple of three coordinates.
You can introduce new names for these compound types (pun unintended). Here, let’s define two new types: Atom and CompoundDict:
| | Atom = tuple[str, tuple[str, str, str]] |
| | CompoundDict = dict[str, Atom] |
They lead to the new function specifications:
| | def read_molecule(reader: TextIO) -> CompoundDict: |
| | """Read a single molecule from reader and return it, or return None to |
| | signal end of file. The returned dictionary has one key/value pair where |
| | the key is the name of the compound and the value is a list of Atoms. |
| | |
| | >>> instring = 'COMPND TEST\\nATOM 1 N 0.1 0.2 0.3\\n'+\ |
| | 'ATOM 2 N 0.2 0.1 0.0\\nEND\\n' |
| | >>> infile = StringIO(instring) |
| | >>> read_molecule(infile) |
| | {'TEST': [('N', ('0.1', '0.2', '0.3')), ('N', ('0.2', '0.1', '0.0'))]} |
| | """ |
| | def read_all_molecules(reader: TextIO) -> CompoundDict: |
| | """Read zero or more molecules from reader, returning a list of the |
| | molecule information. |
| | |
| | >>> cmpnd1 = 'COMPND T1\\nATOM 1 N 0.1 0.2 0.3\\n'+\ |
| | 'ATOM 2 N 0.2 0.1 0.0\\nEND\\n' |
| | >>> cmpnd2 = 'COMPND T2\\nATOM 1 A 0.1 0.2 0.3\\n'+\ |
| | 'ATOM 2 A 0.2 0.1 0.0\\nEND\\n' |
| | >>> infile = StringIO(cmpnd1 + cmpnd2) |
| | >>> result = read_all_molecules(infile) |
| | >>> result['T1'] |
| | [('N', ('0.1', '0.2', '0.3')), ('N', ('0.2', '0.1', '0.0'))] |
| | >>> result['T2'] |
| | [('A', ('0.1', '0.2', '0.3')), ('A', ('0.2', '0.1', '0.0'))] |
| | """ |