Creating New Type Annotations
In addition to type annotations for built-in types such as int and str, module typing provides types List, Set, Tuple, and Dict. For each, you can specify the kind of thing it contains, including type Any for function parameters that work with a mixed set of types.
To explore this, we’ll revisit atoms and molecules from Multiline Records, 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\nATOM 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\nATOM 2 N 0.2 0.1 0.0\nEND\n' |
| | >>> cmpnd2 = 'COMPND T2\nATOM 1 A 0.1 0.2 0.3\nATOM 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 to have read_all_molecules return a dictionary where the keys are the names of compounds and the values are atoms.
Also, instead of using a four-item list for atoms, each atom will be 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, we define new types called Atom and CompoundDict:
| | Atom = Tuple[str, Tuple[str, str, str]] |
| | CompoundDict = Dict[str, Atom] |
That leads to these 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\nATOM 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\nATOM 2 N 0.2 0.1 0.0\nEND\n' |
| | >>> cmpnd2 = 'COMPND T2\nATOM 1 A 0.1 0.2 0.3\nATOM 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'))] |
| | """ |