Merge pull request #543 from Rosalbam1/dev

Multi.py w/ updated syntax

automol/__init__.py

@@ -14,7 +14,7 @@
  - util
  - error
  - mult
- - formula
+ - form
  - inchi_key
  - vmat
  - prop

automol/form/_form.py

@@ -1,13 +1,13 @@
-""" Library for dealing with molecular formula,
-    represented as dict[atom symbol: atom number]
+"""Library for dealing with molecular formula,
+represented as dict[atom symbol: atom number].
 """
 
 import collections
 import functools
 import itertools
-from typing import List
 
 import pyparsing as pp
+from _collections_abc import Sequence
 from pyparsing import pyparsing_common as ppc
 
 from phydat import ptab
@@ -20,14 +20,12 @@
 FORMULA = pp.OneOrMore(ATOM_COUNT)
 
 
-def electron_count(fml):
+def electron_count(fml: dict[str:int]) -> int:
     """Count the number of electrons for the atoms in a molecular formula.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: int
+    :param fml: Stochiometric chemical formula
+    :return: Number of electrons.
     """
-
     assert _is_standard(fml)
 
     elec_count = 0
@@ -38,56 +36,52 @@ def electron_count(fml):
     return elec_count
 
 
-def atom_count(fml):
+def atom_count(fml: dict[str:int]) -> int:
     """Count the number of atoms in this molecular formula.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: int
+    :param fml: Stochiometric chemical formula
+    :return: Number of atoms.
     """
-
     assert _is_standard(fml)
 
     return sum(fml.values())
 
 
-def heavy_atom_count(fml):
+def heavy_atom_count(fml: dict[str:int]) -> int:
     """Count the number of heavy atoms in this molecular formula.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: int
+    :param fml: Stochiometric chemical formula
+    :return: Number of heavy atoms
     """
     assert _is_standard(fml)
     fml = without(fml, ["H"])
     return sum(fml.values())
 
 
-def element_count(fml, symb):
+def element_count(fml: dict[str:int], symb: str) -> int:
     """Count the number of a given element in this molecular formula.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :param symb: atomic symbol of element to be counted
-    :type symb: str
-    :rtype: int
+    :param fml: Stochiometric chemical formula
+    :param symb: Atomic symbol of element to be counted
+    :return: Number of a certain element in the formula
     """
-
     assert _is_standard(fml)
 
     return fml[symb] if symb in fml else 0
 
 
 def without(fml: dict[str, int], symbs: tuple = ()) -> dict[str, int]:
-    """Return a formula without hydrogen
+    """Return a formula without hydrogen.
 
-    :param fml: chemical formula, without hydrogen
+    :param fml: Chemical formula, without hydrogen
+    :symbs: Chemical symbols
+    :return: Dictionary with new formula, without hydrogen
     """
     return {k: v for k, v in fml.items() if k not in symbs}
 
 
 def match(fml1: dict[str, int], fml2: dict[str, int]) -> bool:
-    """Check for a match between two formulas, allowing wildcard values
+    """Check for a match between two formulas, allowing wildcard values.
 
     A stoichiometry of -1 indicates a wildcard value
 
@@ -104,18 +98,14 @@ def match(fml1: dict[str, int], fml2: dict[str, int]) -> bool:
     return fml1 == fml2
 
 
-def add_element(fml, symb, num=1):
-    """add or subtract (if num < 0) this element from the molecular formula
+def add_element(fml: dict[str:int], symb: str, num: int = 1) -> dict[str:int]:
+    """Add or subtract (if num < 0) this element from the molecular formula.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :param symb: atomic symbol of element to be added
-    :type symb: str
-    :param num: number of the element to add to the formula
-    :type num: int
-    :rtype: dict[str:int]
+    :param fml: Stochiometric chemical formula
+    :param symb: Atomic symbol of element to be added
+    :param num: Number of the element to add to the formula
+    :return: Formula with added element
     """
-
     assert ptab.to_number(symb)
     assert _is_standard(fml)
 
@@ -131,57 +121,47 @@ def add_element(fml, symb, num=1):
     return fml
 
 
-def join(fml1, fml2):
+def join(fml1: dict[str:int], fml2: dict[str:int]) -> int:
     """Join two formulas together.
 
-    :param fml1: stochiometric chemical formula 1
-    :type fml1: dict[str:int]
-    :param fml2: stochiometric chemical formula 2
-    :type fml2: dict[str:int]
-    :rtype: int
+    :param fml1: Stochiometric chemical formula 1
+    :param fml2: Stochiometric chemical formula 2
+    :return: Formula with the sum of both formulas
     """
-
     fml = dict(fml1)
     for symb, num in fml2.items():
         fml = add_element(fml, symb, num=num)
 
     return fml
 
 
-def join_sequence(fmls):
-    """Join a sequence of formulas together:
+def join_sequence(fmls: dict[str:int]) -> int:
+    """Join a sequence of formulas together.
 
-    :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: int
+    :param fml: Stochiometric chemical formula
+    :return: Sum of the formulas
     """
-
     return functools.reduce(join, fmls)
 
 
-def sorted_symbols_in_sequence(fmls: List[dict]) -> List[dict]:
-    """Sort a sequence of formulas based on Hill-sorting
+def sorted_symbols_in_sequence(fmls: list[dict]) -> list[dict]:
+    """Sort a sequence of formulas based on Hill-sorting.
 
     :param fmls: A sequence of formulas
-    :type fmls: List[dict]
     :return: The same sequence, but sorted
-    :rtype: List[dict]
     """
     return sorted_symbols(join_sequence(fmls).keys())
 
 
 # Str<->Dict Converters
-def string(fml, hyd=True):
+def string(fml: dict[str:int], hyd: bool = True) -> str:
     """Convert formula dictionary to formula string in the Hill convention.
     Resultant string is identical to InChI formula string.
 
     :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
     :param hyd: include hydrogens?
-    :type hyd: bool
-    :rtype: str
+    :return: True if formula includes hydrogen, False if no hydrogen
     """
-
     fml_lst = [
         (symb, fml[symb]) for symb in sorted_symbols(fml.keys()) if symb != "H" or hyd
     ]
@@ -193,15 +173,13 @@ def string(fml, hyd=True):
     return fml_str
 
 
-def string2(fml):
+def string2(fml: dict[str:int]) -> str:
     """Convert formula dictionary to formula string that includes 1s in when
     there is only one atom.
 
     :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: str
+    :return: Formula which includes 1s when there is only one atom
     """
-
     fml = collections.OrderedDict(sorted(fml.items()))
 
     fml_str = "".join(map(str, itertools.chain.from_iterable(fml.items())))
@@ -223,17 +201,18 @@ def from_string(fml_str: str) -> dict[str, int]:
     return fml
 
 
-def sorted_symbols(seq, symbs_first=("C", "H"), symbs_last=()):
+def sorted_symbols(
+    seq: Sequence,
+    symbs_first: Sequence[str] = ("C", "H"),
+    symbs_last: Sequence[str] = (),
+) -> Sequence:
     """Produce a sorted list of atomic symbols; some elements given priority.
     By default, C placed first, then H, then others in alphabetical order.
 
-    :param seq: formula or sequence of atomic symbols
-    :type seq: dict, list, or tuple
-    :param symbs_first: atomic symbols to place first
-    :type symbs_first: sequence of strings
-    :param symbs_last: atomic symbols to place last
-    :type symbs_last: sequence of strings
-    :rtyp: tuple(str)
+    :param seq: Formula or sequence of atomic symbols
+    :param symbs_first: Atomic symbols to place first
+    :param symbs_last: Atomic symbols to place last
+    :return: Sorted list
     """
 
     def _sort_key(char):
@@ -248,18 +227,19 @@ def _sort_key(char):
     return tuple(sorted(seq, key=_sort_key))
 
 
-def argsort_symbols(seq, symbs_first=("C", "H"), symbs_last=(), idx=None):
+def argsort_symbols(
+    seq: Sequence,
+    symbs_first: Sequence[str] = ("C", "H"),
+    symbs_last: Sequence[str] = (),
+    idx: int | None = None,
+) -> tuple[str, ...]:
     """Determine the sort order for a sequence of atomic symbols.
 
-    :param seq: formula or sequence of atomic symbols
-    :type seq: dict, list, or tuple
-    :param symbs_first: atomic symbols to place first
-    :type symbs_first: sequence of strings
-    :param symbs_last: atomic symbols to place last
-    :type symbs_last: sequence of strings
-    :param idx: index of symbol for sorting
-    :type idx: int
-    :rtype: tuple(int)
+    :param seq: Formula or sequence of atomic symbols
+    :param symbs_first: Atomic symbols to place first
+    :param symbs_last: Atomic symbols to place last
+    :param idx: Index of symbol for sorting
+    :return: Sorted syymboles
     """
 
     def _sort_key(entry):
@@ -287,13 +267,12 @@ def _sort_key(entry):
     )
 
 
-def sort_vector(fml, symbs=None):
-    """Generate a sort vector for sorting various formulas against each other
+def sort_vector(fml: str, symbs: list[str] | None = None):
+    """Generate a sort vector for sorting various formulas against each other.
 
     :param fml_str: stochiometric chemical formula string
-    :type fml_str: str
     :param symbs: atomic symbols in the desired sort order (optional)
-    :type symbs: list[str]
+    :return: Sorted vector
     """
     if symbs is None:
         symbs = sorted_symbols(fml.keys())
@@ -302,14 +281,12 @@ def sort_vector(fml, symbs=None):
     return vec
 
 
-def _is_standard(fml):
+def _is_standard(fml: dict[str:int]) -> bool:
     """Assess if the formula conforms to the standard form.
 
     :param fml: stochiometric chemical formula
-    :type fml: dict[str:int]
-    :rtype: bool
+    :return: True if formula is in standard form, false if not
     """
-
     symbs = list(fml.keys())
 
     return symbs == list(filter(ptab.to_number, map(ptab.to_symbol, symbs)))